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?
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).
