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

Question

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.

Answer 1

The current answers break down if you're using both arguments to construct an error message to pass to the parent Exception class. I believe the best way is to simply override the __reduce__ method in your exception. The __reduce__ method should return a two item tuple. The first item in the tuple is your class. The second item is a tuple containing the arguments to pass to your class's __init__ method.

import pickle

class MyException(Exception):
    def __init__(self, arg1, arg2):
        super().__init__(f'arg1: {arg1}, arg2: {arg2}')
        self.arg1 = arg1
        self.arg2 = arg2
    
    def __reduce__(self):
        return (MyException, (self.arg1, self.arg2))

# Create an exception instance and print info about it
original = MyException('foo', 'bar')
print(repr(original))
print(original.arg1)
print(original.arg2)

# Pickle and unpickle the exception, info printed should match above
reconstituted = pickle.loads(pickle.dumps(original))
print(repr(reconstituted))
print(reconstituted.arg1)
print(reconstituted.arg2)

More info about __reduce__ here.

Answer 2

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'

Answer 3

I like Martijn's answer, but I think a better way is to pass all arguments to the Exception base class:

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

The base Exception class' __reduce__ method will include all the args. By not making all of the extra arguments optional, you can ensure that the exception is constructed correctly.

Answer 4

I simply do this

class MyCustomException(Exception):
    def __init__(self):
        self.value = 'Message about my error'

    def __str__(self):
        return repr(self.value)

... somewhere in code ...
raise MyCustomException

Answer 5

A completely generic solution is to add the following method to your exception class:

    def __reduce__(self):
        return (Exception.__new__, (type(self),) + self.args, self.__dict__)

This works even with keyword only arguments.

Answer 6

I was able to get a custom exception with keyword-only arguments to work with pickling like this:

class CustomException(Exception):
    def __init__(self, arg1, *, kwonly):
        super().__init__(arg1, kwonly) # makes self.args == (arg1, kwonly)
        self.arg1 = arg1
        self.kwonly = kwonly

    def __str__(self):
        # Logic here to turn the args into the real exception message
        return f"{self.arg1} ({self.kwonly})"

    # __repr__ and __eq__ are not needed but are nice for testing
    def __repr__(self):
        return f"{self.__class__.__name__}({self.arg1!r}, kwonly={self.kwonly!r})"

    def __eq__(self, other):
        return isinstance(other, CustomException) and self.args == other.args

    # This is what makes pickling work
    @classmethod
    def _new(cls, arg1, kwonly):
        return cls(arg1, kwonly=kwonly)

    def __reduce__(self):
        return (self._new, (self.arg1, self.kwonly))

Example:

import pickle

w = CustomException('arg1', kwonly='kwonly')
w2 = pickle.loads(pickle.dumps(w))
assert w == w2

By the way, this also works with custom warning classes, since those are also Exception subclasses.

Answer 7

The simplest solution is to pass all of the argument to the base-class's __init__:

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

(Tested with Python versions 3.6, 3.8, 3.9 and 3.10.)