How to pickle a class instance's variables selectively?

Question

I have a class instance:

class A:
    def __init__(self,shape):
        self.shape = shape
        self.matrix = numpy.zeros(shape)

I did some multiplication on self.matrix . When I am pickling the instance of class A , I don't want to pickle self.matrix for some reason.

Currently, my solution is to set self.matrix = None before pickle and reset self.matrix to zeros(shape) after loading from pickled file.

Is there a more elegant solution? Like the transient keyword in Java.

Answer 1

Use the hook methods to limit what is pickled and what is unpickled.

Here, you could use a __getstate__ with corresponding __setstate__ method:

class A:
    def __init__(self, shape):
        self.shape = shape
        self.matrix = numpy.zeros(shape)

    def __getstate__(self):
        return (self.shape,)

    def __setstate__(self, state):
        self.shape, = state
        self.matrix = numpy.zeros(self.shape)

The __setstate__ could be simplified to:

    def __setstate__(self, state):
        self.__init__(*state)

if your __init__ is reusable in this manner. In that case you can also use __getinitargs__ instead of __getstate__ and drop the __setstate__ altogether:

class A:
    def __init__(self, shape):
        self.shape = shape
        self.matrix = numpy.zeros(shape)

    def __getinitargs__(self):
        return (self.shape,)