import numpy as np
import pickle
class RealisticInfoArray(np.ndarray):
def __new__(cls, input_array, info=None):
obj = np.asarray(input_array).view(cls)
obj.info = info
return obj
def __array_finalize__(self, obj):
if obj is None: return
self.info = getattr(obj, 'info', None)
def __reduce__(self):
# Get the parent's __reduce__ tuple
pickled_state = super(RealisticInfoArray, self).__reduce__()
# Create our own tuple to pass to __setstate__
new_state = pickled_state[2] + (self.__dict__,)
# Return a tuple that replaces the parent's __setstate__ tuple with our own
return (pickled_state[0], pickled_state[1], new_state)
def __setstate__(self, state):
self.__dict__.update(state[-1]) # Update the internal dict from state
# Call the parent's __setstate__ with the other tuple elements.
super(RealisticInfoArray, self).__setstate__(state[0:-1])
class VeryRealisticInfoArray(RealisticInfoArray):
def __new__(cls, input_array, info, more_info):
obj = super().__new__(cls, input_array, info)
obj.more_info = more_info
return obj
r = VeryRealisticInfoArray(np.arange(19), info=42, more_info=1234)
print(r.shape, r.info, r.more_info)
rr = pickle.loads(pickle.dumps(r))
print(rr.shape, rr.info, rr.more_info)
