Python - Pass class as argument to decorator of method within that class?

user_2681057 0
Viewed 2001

I'm looking to do the following:

class A(object):
    def first_method(self):
        print "I'm the first method!"    

    @some_decorator(A.first_method)
    def second_method(self):
        print "I'm the second method!"

But I'm running into the problem that A is undefined within itself at the time that the decorator is parsed. Is there any way for me to reference A in the decorator? Alternatively, if I just pass the decorator the bound method first_method is it possible for me to recover that first_method belongs to A?

decorator python
2 Answers

In python 3, you can just say @some_decorator(first_method), and it will work, as all methods are normal functions in the classes as the containers.

In python 2, there is a complicated system of bound & unbound, instance-, class-, and static methods. Due to this, you cannot access first_method while inside of a class definitionm (until the class if fully formed).

A little workaround would be to split that class into two classes:

class BaseA(object):
    def first_method(self):
        print "I'm the first method!"    

class A(BaseA):
    @some_decorator(BaseA.first_method)
    def second_method(self):
        print "I'm the second method!"

Not the best solution for all cases, but will work.

Also, keep in mind, that in both cases (py2 & py3), the decorator will refer to the first_method as it was declared here. If any descendant class redefines the method, the new method will NOT be used in the decorator; only the parent one will be used.

Probably, you should not refer to first_method at all. Instead, just accept self/cls first positional argument in the decorator's wrapper, and use self.first_method/cls.first_method there:

import functools

def some_decorator(fn):
    @functools.wraps(fn)
    def wrapper(self, *args, **kwargs):
        first_method = self.first_method
        first_method()
        return fn(self, *args, **kwargs)
    return wrapper

class A(object):
    def first_method(self):
        print "I'm the first method of A!"    

    @some_decorator
    def second_method(self):
        print "I'm the second method!"


class B(A):
    def first_method(self):
        print "I'm the first method of B!"    


A().second_method()
# I'm the first method of A!
# I'm the second method!

B().second_method()
# I'm the first method of B!
# I'm the second method!

If you want to make that method configurable:

def some_decorator(method_name):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(self, *args, **kwargs):
            first_method = getattr(self, method_name)
            first_method()
            return fn(self, *args, **kwargs)
        return wrapper
    return decorator


class A(object):
    def first_method(self):
        print "I'm the first method of A!"    

    @some_decorator('first_method')
    def second_method(self):
        print "I'm the second method!"

class B(A):
    def first_method(self):
        print "I'm the first method of B!"
user_8573830

You can use your decorator as a classic function (without the @):

def some_decorator(arg):
    # ...
    pass


class A(object):
    def first_method(self):
        print("I'm the first method!")

    def second_method(self):
        print("I'm the second method!")


A.second_method = some_decorator(A.first_method)(A.second_method)
user_15139330
Related