Class decorator for methods from other class

Question

NOTE:
I've got a related question here: How to access variables from a Class Decorator from within the method it's applied on?

---

I'm planning to write a fairly complicated decorator. Therefore, the decorator itself should be a class of its own. I know this is possible in Python (Python 3.8):

import functools

class MyDecoratorClass:
    def __init__(self, func):
        functools.update_wrapper(self, func)
        self.func = func
    
    def __call__(self, *args, **kwargs):
        # do stuff before
        retval = self.func(*args, **kwargs)
        # do stuff after
        return retval

@MyDecoratorClass
def foo():
    print("foo")

Now my problem starts when I try to apply the decorator on a method instead of just a function - especially if it's a method from another class. Let me show you what I've tried:

 

1. Trial one: identity loss

The decorator MyDecoratorClass below doesn't (or shouldn't) do anything. It's just boilerplate code, ready to be put to use later on. The method foo() from class Foobar prints the object it is called on:

import functools

class MyDecoratorClass:
    def __init__(self, method):
        functools.update_wrapper(self, method)
        self.method = method

    def __call__(self, *args, **kwargs):
        # do stuff before
        retval = self.method(self, *args, **kwargs)
        # do stuff after
        return retval

class Foobar:
    def __init__(self):
        # initialize stuff
        pass

    @MyDecoratorClass
    def foo(self):
        print(f"foo() called on object {self}")
        return

Now what you observe here is that the self in the foo() method gets swapped. It's no longer a Foobar() instance, but a MyDecoratorClass() instance instead:

>>> foobar = Foobar()
>>> foobar.foo()
foo() called from object <__main__.MyDecoratorClass object at 0x000002DAE0B77A60>

In other words, the method foo() loses its original identity. That brings us to the next trial.

 

2. Trial two: keep identity, but crash

I attempt to preserve the original identity of the foo() method:

import functools

class MyDecoratorClass:
    def __init__(self, method):
        functools.update_wrapper(self, method)
        self.method = method

    def __call__(self, *args, **kwargs):
        # do stuff before
        retval = self.method(self.method.__self__, *args, **kwargs)
        # do stuff after
        return retval

class Foobar:
    def __init__(self):
        # initialize stuff
        pass

    @MyDecoratorClass
    def foo(self):
        print(f"foo() called on object {self}")
        return

Now let's test:

>>> foobar = Foobar()
>>> foobar.foo()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 7, in __call__
AttributeError: 'function' object has no attribute '__self__'

Yikes!

---

EDIT
Thank you @AlexHall and @juanpa.arrivillaga for your solutions. They both work. However, there is a subtle difference between them.

Let's first take a look at this one:

def __get__(self, obj, objtype) -> object:
    temp = type(self)(self.method.__get__(obj, objtype))
    print(temp)
    return temp

I've introduced a temporary variable, just to print what __get__() returns. Each time you access the method foo(), this __get__() function returns a new MyDecoratorClass() instance:

>>> f = Foobar()
>>> func1 = f.foo
>>> func2 = f.foo
>>> print(func1 == func2)
>>> print(func1 is func2)
<__main__.MyDecoratorClass object at 0x000001B7E974D3A0>
<__main__.MyDecoratorClass object at 0x000001B7E96C5520>
False
False

The second approach (from @juanpa.arrivillaga) is different:

def __get__(self, obj, objtype) -> object:
    temp = types.MethodType(self, obj)
    print(temp)
    return temp

The output:

>>> f = Foobar()
>>> func1 = f.foo
>>> func2 = f.foo
>>> print(func1 == func2)
>>> print(func1 is func2)
<bound method Foobar.foo of <__main__.Foobar object at 0x000002824BBEF4C0>>
<bound method Foobar.foo of <__main__.Foobar object at 0x000002824BBEF4C0>>
True
False

There is a subtle difference, but I'm not sure why.

Answer 1

Functions are descriptors and that's what allows them to auto-bind self. The easiest way to deal with this is to implement decorators using functions so that this is handled for you. Otherwise you need to explicitly invoke the descriptor. Here's one way:

import functools


class MyDecoratorClass:
    def __init__(self, method):
        functools.update_wrapper(self, method)
        self.method = method

    def __get__(self, instance, owner):
        return type(self)(self.method.__get__(instance, owner))

    def __call__(self, *args, **kwargs):
        # do stuff before
        retval = self.method(*args, **kwargs)
        # do stuff after
        return retval


class Foobar:
    def __init__(self):
        # initialize stuff
        pass

    @MyDecoratorClass
    def foo(self, x, y):
        print(f"{[self, x, y]=}")


@MyDecoratorClass
def bar(spam):
    print(f"{[spam]=}")


Foobar().foo(1, 2)
bar(3)

Here the __get__ method creates a new instance of MyDecoratorClass with the bound method (previously self.method was just a function since no instance existed yet). Also note that __call__ just calls self.method(*args, **kwargs) - if self.method is now a bound method, the self of FooBar is already implied.

Answer 2

You can implement the descriptor protocol, an example of how functions do it (but in pure python) is available in the Descriptor HOWTO, translated to your case:

import functools
import types

class MyDecoratorClass:
    def __init__(self, func):
        functools.update_wrapper(self, func)
        self.func = func

    def __call__(self, *args, **kwargs):
        # do stuff before
        retval = self.func(*args, **kwargs)
        # do stuff after
        return retval

    def __get__(self, obj, objtype=None):
        if obj is None:
            return self
        return types.MethodType(self, obj)

Note, return types.MethodType(self, obj) is essentially equivalent to

return lambda *args, **kwargs : self.func(obj, *args, **kwargs)

---

Note from Kristof
Could it be that you meant this:

return types.MethodType(self, obj) is essentially equivalent to

return lambda *args, **kwargs : self(obj, *args, **kwargs)

Note that I replaced self.func(..) with self(..). I tried, and only this way I can ensure that the statements at # do stuff before and # do stuff after actually run.