Search code examples
pythonclassobjectpickle

How to update class attributes with load function


I have a class in which the method save stores the created object to disk using the pickle package:

def save(
        self,
        filepath,
        ):
    #CHECK ATTRIBUTE VALUES
    pickle_out = open(filepath+'.pickle', 'wb')
    pickle.dump(self, pickle_out)
    pickle_out.close()

Similarly, I want to add a load method that loads a pickled file and updates the object with the loaded attributes.

In my load method, I used:

infile = open(filepath,'rb')
self = pickle.load(infile)
infile.close()

However, this does not update the class attributes. How can I write a load method that does update the class variables?


Solution

  • When not using a C extension or slots, class instance data is stored in a dictionary. You could unpickle the class and replace that dictionary.

    import pickle
    
    class Foo:
    
        def __init__(self, val):
            self.val = val
    
        def save(self, filepath):
            with open(filepath,'wb') as outfile:
                pickle.dump(self, outfile)
    
        def load(self, filepath):
            with open(filepath,'rb') as infile:
                self.__dict__ = pickle.load(infile).__dict__
    
    f1 = Foo(1)
    f1.save("test.pkl")
    del f1
    f2 = Foo(2)
    print(f2.val)
    f2.load("test.pkl")
    print(f2.val)
    

    It would make sense to pickle only self.__dict__ because that is the unique data of interest.

    I don't see any benefit to this over the normal way of using pickle - have regular function or staticmethod load and return the pickled object. That is the intended use of pickle.