python - How to pickle and unpickle instances of a class that inherits from defaultdict?

Question

I have a class that inherits from defaultdict like this:

class listdict(defaultdict):
    def __init__(self):
        defaultdict.__init__(self, list)

I can pickle it, but when I unpickle it, this happens:

('__init__() takes exactly 1 argument (2 given)', <class 'listdict'>, (<type 'list'>,))

The class does not define any special methods of the pickle protocol. Pickling and unpickling of a normal defaultdict(list) works as expected. Can anyone enlighten me?

Answer 1

Types define how instances of it get pickled by defining one or more of a (fairly large) set of methods. Each has its own subtle behaviour. See the docs on the pickle protocol. In the case of collections.defaultdict, it uses the __reduce__ method:

>>> l = collections.defaultdict(list)
>>> l.__reduce__()
(<type 'collections.defaultdict'>, (<type 'list'>,), None, None, <dictionary-itemiterator object at 0x7f031fb3c470>)

The first item in the tuple there is the type, and the second item is the tuple of arguments to pass to the type when instantiating it. If you don't override __reduce__, the first item will correctly change to your type, but the second item will not. This causes the error you see. A crude example of how you could fix it:

>>> import collections
>>> import pickle
>>> class C(collections.defaultdict):
...     def __init__(self):
...         collections.defaultdict.__init__(self, list)
...     def __reduce__(self):
...         t = collections.defaultdict.__reduce__(self)
...         return (t[0], ()) + t[2:]
...
>>> c = C()
>>> c[1].append(2)
>>> c[2].append(3)
>>> c2 = pickle.loads(pickle.dumps(c))
>>> c2 == c
True

It's only a crude example because there's more to pickling (like __reduce_ex__) and it's all fairly intricate. In this case, using __getinitargs__ may be more convenient.

Alternatively, you could make your class's __init__ method take an optional callable, defaulting to list, or you could just use a function instead of a class:

def listdict():
    return collections.defaultdict(list)