Torxed
Torxed

Reputation: 23480

pickling dict inherited class is missing internal values

I've played around for a bit with the code and obviously the reason for the failure is that when setting 'wham' the value is another instance of the class TestDict which works fine as long as i don't try to pickle and unpickle it.
Because if i do self.test is missing.

Traceback:

Traceback (most recent call last):
  File "test.py", line 30, in <module>
    loads_a = loads(dumps_a)
  File "test.py", line 15, in __setitem__
    if self.test == False:
AttributeError: 'TestDict' object has no attribute 'test'

The code:

from pickle import dumps, loads

class TestDict(dict):
    def __init__(self, test=False, data={}):
        super().__init__(data)
        self.test = test

    def __getitem__(self, k):
        if self.test == False:
            pass
        return dict.__getitem__(self, k)

    def __setitem__(self, k, v):
        if self.test == False:
            pass
        if type(v) == dict:
            super().__setitem__(k, TestDict(False, v))
        else:
            super().__setitem__(k, v)

if __name__ == '__main__':
    a = TestDict()
    a['wham'] = {'bam' : 1}
    b = TestDict(True)
    b['wham'] = {'bam' : 2}
    dumps_a = dumps(a)
    dumps_b = dumps(b)

    loads_a = loads(dumps_a)
    loads_b = loads(dumps_b)

    print(loads_a)
    print(loads_b)

The code works if not replacing __setitem__ and __getitem__ but i want to add extra functionality to those two specific functions.

I've also tried:

class TestDict(dict):
    __module__ = os.path.splitext(os.path.basename(__file__))[0]

Which sort of worked, as long as i don't nest TestDict within TestDict meaning i won't get to replace __setitem__ and __getitem__ in sub-parts of the dictionary.

Upvotes: 2

Views: 817

Answers (2)

falsetru
falsetru

Reputation: 369134

Define __reduce__ method:

class TestDict(dict):
    ...

    def __reduce__(self):
        return type(self), (self.test, dict(self))

Upvotes: 3

Viktor Kerkez
Viktor Kerkez

Reputation: 46596

The problem is that pickle doesn't call the __init__ method of the object when it does the unpicling, so the self.test variable is not created at the moment when it tries to set the items of the dictionary.

Apparently setting the attributes is staged after setting items of the dictionary.

One way to solve it is to add a class level attribute that will be overridden in the instances:

class TestDict(dict):
    test = False
    ...

Upvotes: 1

Related Questions