Search code examples
pythonexceptionserialization

How to make a custom exception class with multiple init args pickleable


Why does my custom Exception class below not serialize/unserialize correctly using the pickle module?

import pickle

class MyException(Exception):
    def __init__(self, arg1, arg2):
        self.arg1 = arg1
        self.arg2 = arg2

        super(MyException, self).__init__(arg1)

e = MyException("foo", "bar")

str = pickle.dumps(e)
obj = pickle.loads(str)

This code throws the following error:

Traceback (most recent call last):
File "test.py", line 13, in <module>
   obj = pickle.loads(str)
File "/usr/lib/python2.7/pickle.py", line 1382, in loads
   return Unpickler(file).load()
File "/usr/lib/python2.7/pickle.py", line 858, in load
   dispatch[key](self)
File "/usr/lib/python2.7/pickle.py", line 1133, in load_reduce
   value = func(*args)
TypeError: __init__() takes exactly 3 arguments (2 given)

I'm sure this problem stems from a lack of knowledge on my part of how to make a class pickle-friendly. Interestingly, this problem doesn't occur when my class doesn't extend Exception.


Solution

  • Make arg2 optional:

    class MyException(Exception):
        def __init__(self, arg1, arg2=None):
            self.arg1 = arg1
            self.arg2 = arg2
            super(MyException, self).__init__(arg1)
    

    The base Exception class defines a .__reduce__() method to make the extension (C-based) type picklable and that method only expects one argument (which is .args); see the BaseException_reduce() function in the C source.

    The easiest work-around is making extra arguments optional. The __reduce__ method also includes any additional object attributes beyond .args and .message and your instances are recreated properly:

    >>> e = MyException('foo', 'bar')
    >>> e.__reduce__()
    (<class '__main__.MyException'>, ('foo',), {'arg1': 'foo', 'arg2': 'bar'})
    >>> pickle.loads(pickle.dumps(e))
    MyException('foo',)
    >>> e2 = pickle.loads(pickle.dumps(e))
    >>> e2.arg1
    'foo'
    >>> e2.arg2
    'bar'