Skip to content

Instantly share code, notes, and snippets.

@cjw85
Created January 18, 2016 21:34
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save cjw85/c9ecae0114430e235b18 to your computer and use it in GitHub Desktop.
Save cjw85/c9ecae0114430e235b18 to your computer and use it in GitHub Desktop.
Nicely pickling numpy array
import numpy as np
import pickle
class MyArray(np.ndarray):
"""An np.ndarray subclass with attrs that behave nicely."""
def __new__(cls, a, meta=None):
obj = np.asarray(a).view(cls)
obj.meta = meta
return obj
def __array_finalize__(self, obj):
if obj is None: return
self.meta = getattr(obj, 'meta', None)
if hasattr(obj, '__dict__'):
for attr, val in obj.__dict__.items():
setattr(self, attr, val)
def __array_wrap__(self, out_arr, context=None):
return np.ndarray.__array_wrap__(self, out_arr, context)
def __reduce__(self):
cons, args, state = super(MyArray, self).__reduce__()
state = (state,) + (self.__dict__,)
return (cons, args, state)
def __setstate__(self, d):
standard, extra = d
super(MyArray, self).__setstate__(standard)
for attr, val in extra.items():
setattr(self, attr, val)
# An array that has an attribute
a = MyArray(np.array(range(5)), meta='meta_data')
# Can patch any other attributes
a.info = 'info_data'
# Pickling attributes works
b = pickle.loads(pickle.dumps(a))
print b, b.info, b.meta
# Attributes copied when ufunc applied
c = b + 1
print c, c.info, c.meta
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment