3.8.1 Example

Here's a simple example of how to modify pickling behavior for a class. The TextReader class opens a text file, and returns the line number and line contents each time its readline() method is called. If a TextReader instance is pickled, all attributes except the file object member are saved. When the instance is unpickled, the file is reopened, and reading resumes from the last location. The __setstate__() and __getstate__() methods are used to implement this behavior.

# illustrate __setstate__ and __getstate__  methods
# used in pickling.

class TextReader:
    "Print and number lines in a text file."
    def __init__(self,file):
        self.file = file
        self.fh = open(file,'r')
        self.lineno = 0

    def readline(self):
        self.lineno = self.lineno + 1
        line = self.fh.readline()
        if not line:
            return None
        return "%d: %s" % (self.lineno,line[:-1])

    # return data representation for pickled object
    def __getstate__(self):
        odict = self.__dict__    # get attribute dictionary
        del odict['fh']          # remove filehandle entry
        return odict

    # restore object state from data representation generated 
    # by __getstate__
    def __setstate__(self,dict):
        fh = open(dict['file'])  # reopen file
        count = dict['lineno']   # read from file...
        while count:             # until line count is restored
            fh.readline()
            count = count - 1
        dict['fh'] = fh          # create filehandle entry
        self.__dict__ = dict     # make dict our attribute dictionary

A sample usage might be something like this:

>>> import TextReader
>>> obj = TextReader.TextReader("TextReader.py")
>>> obj.readline()
'1: #!/usr/local/bin/python'
>>> # (more invocations of obj.readline() here)
... obj.readline()
'7: class TextReader:'
>>> import pickle
>>> pickle.dump(obj,open('save.p','w'))

  (start another Python session)

>>> import pickle
>>> reader = pickle.load(open('save.p'))
>>> reader.readline()
'8:     "Print and number lines in a text file."'


Send comments on this document to python-docs@python.org.