Pickle and decorated classes (PicklingError: not the same object)

Question

The following minimal example uses a dummy decorator, that justs prints some message when an object of the decorated class is constructed.

import pickle


def decorate(message):
    def call_decorator(func):
        def wrapper(*args, **kwargs):
            print(message)
            return func(*args, **kwargs)

        return wrapper

    return call_decorator


@decorate('hi')
class Foo:
    pass


foo = Foo()
dump = pickle.dumps(foo) # Fails already here.
foo = pickle.loads(dump)

Using it however makes pickle raise the following exception:

_pickle.PicklingError: Can't pickle <class '__main__.Foo'>: it's not the same object as __main__.Foo

Is there anything I can do to fix this?

Ethelyn

__main__.Foo has been replaced by wrapper function, so it is no longer a class. Pickle can't handle that case, because the foo.__class__ attribute points to a class object that pickle can't load. What is the goal of the decorator? – Cowie

@MartijnPieters: It's the type_checked_call decorator in my library: github.com/Dobiasd/undictify – Ethelyn

That decorator should return the class unchanged. Replace the __init__ or __new__ method of the class with a wrapper, instead. – Cowie

Answer 1

Pickle requires that the __class__ attribute of instances can be loaded via importing.

Pickling instances only stores the instance data, and the __qualname__ and __module__ attributes of the class are used to later on re-create the instance by importing the class again and creating a new instance for the class.

Pickle validates that the class can actually be imported first. The __module__ and __qualname__ pair are used to find the correct module and then access the object named by __qualname__ on that module, and if the __class__ object and the object found on the module don't match, the error you see is raised.

Here, foo.__class__ points to a class object with __qualname__ set to 'Foo' and __module__ set to '__main__', but sys.modules['__main__'].Foo doesn't point to a class, it points to a function instead, the wrapper nested function your decorator returned.

There are two possible solutions:

• Don't return a function, return the original class, and perhaps instrument the class object to do the work the wrapper does. If you are acting on the arguments for the class constructor, add or wrap a __new__ or __init__ method on the decorated class.

  Take into account that unpickling usually calls __new__ on the class to create a new empty instance, before restoring the instance state (unless pickling has been customised).

• Store the class under a new location. Alter the __qualname__ and perhaps the __module__ attributes of the class to point to a location where the original class can be found by pickle. On unpickling the right type of instance will be created again, just like the original Foo() call would have.

Another option is to customise pickling for the produced class. You can give the class new __reduce_ex__ and new __reduce__ methods that point to the wrapper function or a custom reduce function, instead. This can get complex, as the class may already have customised pickling, and object.__reduce_ex__ provides a default, and the return value can differ by pickle version.

If you don't want to alter the class, you can also use the copyreg.pickle() function to register a custom __reduce__ handler for the class.

Either way, the return value of the reducer should still avoid referencing the class and should reference the new constructor instead, by the name that it can be imported with. This can be problematic if you use the decorator directly with new_name = decorator()(classobj). Pickle itself would not deal with such situations either (as classobj.__name__ would not match newname).

Cowie

Thanks a lot. I'm now replacing the __new__ function of the class in my updated minimal example. It's OK in this use case, but a problem arises when the decorator function is not used as a decorator, but called normally instead. In that case other parts of the code should not be affected, but since we patch the whole class, it actually does (code). Any idea how this could be solved? – Ethelyn

@TobiasHermann: yes, if you modify the class in place in multiple you can't re-use the decorator to produce multiple 'copies'. Another option is to subclass and give the subclass the same name and module, perhaps. Subclassing requires that the classes themselves behave correctly when it comes to calling overridden methods, however, as the Foo.somemethod reference would now pass through your shim subclass depending on if super() was used or a direct reference. The MRO of subclasses may be affected too. – Cowie

You mean like this? (It seems to work.) – Ethelyn

@TobiasHermann: note that any other use of newname = decorator()(classobj) would also fail to work with pickle anyway, as pickle demands that classobj.__name__ be available for import. – Cowie

@TobiasHermann: your example doesn't use cooperative inheritance. Try subclassing the decorated result with mixin classes or a diamond-inheritance, and __init__ methods using super().__init__() to call the next initialiser in MRO order. You hardcode the next object whose __init__ method is called by using the_class.__init__. And if you switched to super(self, the_class).__init__(*args, **kwargs) instead, you'll break code that uses Foo.__init__() calls in a subclass (and that code would break today, as Foo.__init__ can't be reached through your wrapper). – Cowie

@TobiasHermann: and that's the other big problem with your library, you replace a class with a function, so all ClassName.attributename access is now broken. – Cowie

You are right. We get a AttributeError: Can't pickle local object 'decorate.<locals>.call_decorator.<locals>.wrapper'. So it's not perfect but maybe better than the original version, which was unable to pickle at all. The new version can at least pickle objects created from decorated classes. - The "you replace a class with a function" seems to be about to be fixed thanks to your help. ;) - Your suggestion with "mixin classes or a diamond-inheritance" is something I don't yet understand. But of course I'll try. – Ethelyn

That's because you didn't adjust the __qualname__ of the new class to match the old class. – Cowie

You mean like so? – Ethelyn

As for super() cooperative examples, see rhettinger.wordpress.com/2011/05/26/super-considered-super or the PyCon talk by the same author on the subject. – Cowie

@TobiasHermann: your last gist fails again with _pickle.PicklingError: Can't pickle <class '__main__.Foo'>: it's not the same object as __main__.Foo because you are not pickling Foo, you are pickling the foo1 class. And that kind of aliasing and swapping out wouldn't work without decorating either (class Foo: pass, then Foo1 = Foo, and class Foo(Foo1): pass would fail to pickle Foo1() instances). – Cowie

Yes, I know it fails. Sorry, I did not make this clear. :) - I'll try to dive deeper into the whole thing, but currently the version with replacing __new__ seems quite good to me. I will simply not support calling the decorator normally with a class. For that case foo = decorate('hi')(Foo.__new__)(Foo) can be used anyways, like so. – Ethelyn

In case you are interested, my library now has a better decorator (@type_checked_constructor) that only replaces the constructor of a class with a type-checking wrapper, thus no longer destroying other important information of the class by replacing it completely with a function like before. So, thanks again very much for your help. :) – Ethelyn

If I understood well, to sum up your first four paragraphs, pickle.dumps(obj) checks that getattr(sys.modules[type(obj).__module__], type(obj).__qualname__) == type(obj). – Bold

Do you know why we get the same _pickle.PicklingError when instantiating a multiprocessing.Manager() after @georgexsh's Python workaround multiprocessing.reduction.register(types.MethodType, my_reduce) here? – Bold

@Maggyero: sorry, I don't off-hand and I don't have much time at the moment to chase down why. – Cowie

Answer 2

Using dill, istead of pickle raises no errors.

import dill


def decorate(message):
    def call_decorator(func):
        def wrapper(*args, **kwargs):
            print(message)
            return func(*args, **kwargs)

        return wrapper

    return call_decorator


@decorate('hi')
class Foo:
    pass


foo = Foo()
dump = dill.dumps(foo) # Fails already here.
foo = dill.loads(dump)

output -> hi

Northwestwards