CODEWITHSUNDEEP
pythonpython-3.xpicklepython-decorators
rohit-biswas
rohit-biswas

Reputation: 845

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'

Upvotes: 6

Views: 1553

Answers (1)

user2357112
user2357112

Reputation: 281683

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

Upvotes: 9

Related Questions