Art
Art

Reputation: 24577

Exclude object's field from pickling in python

I would like to avoid pickling of certain fields in an instance of a class. Currently, before pickling I just set those fields to None, but I wonder whether there's more elegant solution?

Upvotes: 48

Views: 29316

Answers (4)

Collin Heist
Collin Heist

Reputation: 2682

For those whose classes define __slots__, and therefore do not have __dict__ for use as suggested in the current answer, the following solution will work:

class Foo:
    __slots__ = ('a', 'b', 'bad_attr')

    def __init__(self) -> None:
        self.a = 1
        self.b = 2
        self.bad_attr = 'foo'

    def __getstate__(self) -> dict:
        return {
            attr: getattr(self, attr) for attr in self.__slots__
            if attr not in ('bad_attr', ) # Attributes to exclude
        }

    def __setstate__(self, state: dict) -> None:
        for attr, value in state.items():
            setattr(self, attr, value)

Upvotes: 0

anijhaw
anijhaw

Reputation: 9402

One way to handle instance attributes that are not picklable objects is to use the special methods available for modifying a class instance's state: __getstate__() and __setstate__(). Here is an example

class Foo(object):

    def __init__(self, value, filename):
        self.value = value
        self.logfile = file(filename, 'w')

    def __getstate__(self):
        """Return state values to be pickled."""
        f = self.logfile
        return (self.value, f.name, f.tell())

    def __setstate__(self, state):
        """Restore state from the unpickled state values."""
        self.value, name, position = state
        f = file(name, 'w')
        f.seek(position)
        self.logfile = f

When an instance of Foo is pickled, Python will pickle only the values returned to it when it calls the instance's __getstate__() method. Likewise, during unpickling, Python will supply the unpickled values as an argument to the instance's __setstate__() method. Inside the __setstate__() method we are able to recreate the file object based on the name and position information we pickled, and assign the file object to the instance's logfile attribute.

Reference: http://www.ibm.com/developerworks/library/l-pypers.html

Upvotes: 58

congusbongus
congusbongus

Reputation: 14622

Pickling uses the object's __getstate__ and __setstate__ methods; you can override them and ignore the fields you want.

# foo.py
class Foo:
    def __init__(self):
        self.bar = 1
        self.baz = 2

    def __getstate__(self):
        state = self.__dict__.copy()
        # Don't pickle baz
        del state["baz"]
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        # Add baz back since it doesn't exist in the pickle
        self.baz = 0
# main.py
import pickle

from foo import Foo


foo = Foo()
print(f"Foo bar: {foo.bar} baz: {foo.baz}")

new_foo = pickle.loads(pickle.dumps(foo))
print(f"New bar: {new_foo.bar} baz: {new_foo.baz}")

Output:

Foo bar: 1 baz: 2
New bar: 1 baz: 0

You can find another example here: https://docs.python.org/3/library/pickle.html#handling-stateful-objects

Upvotes: 32

unutbu
unutbu

Reputation: 879729

There is an example in the docs which solves your problem with __getstate__ and __setstate__.

Upvotes: 26

Related Questions