Accessing bound method or self when decorating a method

Question

I have a use case where I want to decorate a method with an additional way to call it, such as in this code:

def decorator(func):
    def enhanced(*args, **kwargs):
        func(*args, **kwargs)

    func.enhanced = enhanced
    return func

@decorator
def function():
    pass

class X:
    @decorator
    def function(self):
        pass

x = X()

function()
function.enhanced()
x.function()
# x.function.enhanced()
x.function.enhanced(x)

The first three calls work as expected, but x.function.enhanced() does not; I'd have to write x.function.enhanced(x) to make it work. I know that this is because the func passed to the decorator is not a bound method but a function, and thus needs to be passed self explicitly.

But how can I get around this? From the little bit I understand about descriptors, they're only relevant when looking up on a class, and as func is not a class, func.enhanced is not looked up in a way that I can intercept.

Is there something I can do here?

Answer 1

You can return a descriptor that returns an object that makes itself callable and has an enhanced attribute mapped to your enhanced wrapper function:

from functools import partial
def decorator(func):
    class EnhancedProperty:
        # this allows function.enhanced() to work
        def enhanced(self, *args, **kwargs):
            print('enhanced', end=' ') # this output is for the demo below only
            return func(*args, **kwargs)
        # this allows function() to work
        def __call__(self, *args, **kwargs):
            return func(*args, **kwargs)
        def __get__(self, obj, objtype):
            class Enhanced:
                # this allows x.function() to work
                __call__ = partial(func, obj)
                # this allows x.function.enhanced() to work
                enhanced = partial(self.enhanced, obj)
            return Enhanced()
    return EnhancedProperty()

so that:

@decorator
def function():
    print('function')

class X:
    @decorator
    def function(self):
        print('method of %s' % self.__class__.__name__)

x = X()

function()
function.enhanced()
x.function()
x.function.enhanced()

would output:

function
enhanced function
method of X
enhanced method of X