pythonpython-3.xpicklepython-decorators

Extending a class in Python inside a decorator


I am using a decorator to extend certain classes and add some functionality to them, something like the following:

def useful_stuff(cls):
  class LocalClass(cls):
    def better_foo(self):
      print('better foo')
  return LocalClass

@useful_stuff
class MyClass:
  def foo(self):
    print('foo')

Unfortunaltely, MyClass is no longer pickleable due to the non global LocalClass

AttributeError: Can't pickle local object 'useful_stuff.<locals>.LocalClass'

Solution

  • You need to set the metadata so the subclass looks like the original:

    def deco(cls):
        class SubClass(cls):
            ...
        SubClass.__name__ = cls.__name__
        SubClass.__qualname__ = cls.__qualname__
        SubClass.__module__ = cls.__module__
        return SubClass
    

    Classes are pickled by using their module and qualname to record where to find the class. Your class needs to be found in the same location the original class would have been if it hadn't been decorated, so pickle needs to see the same module and qualname. This is similar to what functools.wraps does for decorated functions.

    However, it would probably be simpler and less bug-prone to instead add the new methods directly to the original class instead of creating a subclass:

    def better_foo(self):
        print('better_foo')
    
    def useful_stuff(cls):
        cls.better_foo = better_foo
        return cls