Extending a class in Python inside a decorator

Question

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'

• I need to pickle my classes. Can you recommend a better design?
• Considering that there can be multiple decorators on a class, would switching to multiple inheritance by having MyClass inherit all the functionality be a better option?

Answer 1

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 funcutils.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