Why can't instances of this subclass of set be unpickled?

user_8702322 0
Viewed 59

When I dump a class instance to a file with pickle module, it's successful, however, when I load the instance from the file, it failed with TypeError exception raised.

I found all subclass that derived from set will raise a TypeError when attempting to unpickle it.

class TreeKeys(set):
    def __init__(self):
        super(TreeKeys, self).__init__()

    def add(self, tk):
        assert tk.__class__ == tuple
        super(TreeKeys, self).add(tk)


if __name__ == '__main__':
    a = TreeKeys()
    a.add((1,2,3))
    with open('tmp.pickle', 'wb') as tmp_fd:
        pickle.dump(a, tmp_fd)
    with open('tmp.pickle', 'rb') as tmp_fd:
        obj = pickle.load(tmp_fd)  # this is the code line raise the TypeError exception.
    pass

Connected to pydev debugger (build 182.3684.100)
Traceback (most recent call last):
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2018.2\helpers\pydev\pydevd.py", line 1664, in <module>
    main()
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2018.2\helpers\pydev\pydevd.py", line 1658, in main
    globals = debugger.run(setup['file'], None, None, is_module)
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2018.2\helpers\pydev\pydevd.py", line 1068, in run
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "G:/pycharm-projects/new_keyinfo/verify_treekeys.py", line 56, in <module>
    obj = pickle.load(tmp_fd)
  File "C:\Python27_64bit\Lib\pickle.py", line 1384, in load
    return Unpickler(file).load()
  File "C:\Python27_64bit\Lib\pickle.py", line 864, in load
    dispatch[key](self)
  File "C:\Python27_64bit\Lib\pickle.py", line 1139, in load_reduce
    value = func(*args)
TypeError: __init__() takes exactly 1 argument (2 given)

Process finished with exit code -1
pickle python-2.7 set python
1 Answers

The problem is you didn't define your class' __init__() method to take an (optional) iterable argument like its base class, the built-in set class does, and pickle tries to call it with one to restore the instance's contents when it loads the file.

Here's a simple way to fix that:

import pickle

class TreeKeys(set):
    def __init__(self, iterable=None):
        if iterable is None:
            super(TreeKeys, self).__init__()
        else:
            super(TreeKeys, self).__init__(iterable)

    def add(self, tk):
        assert tk.__class__ == tuple
        super(TreeKeys, self).add(tk)


if __name__ == '__main__':
    a = TreeKeys()
    a.add((1,2,3))
    with open('tmp.pickle', 'wb') as tmp_fd:
        pickle.dump(a, tmp_fd)
    with open('tmp.pickle', 'rb') as tmp_fd:
        obj = pickle.load(tmp_fd)  # This will now work.
    print(obj)  # -> TreeKeys([(1, 2, 3)])
user_3552300
Related