Python Decorators in Classes

Python decorators in classes

Would something like this do what you need?

class Test(object):
    def _decorator(foo):
        def magic( self ) :
            print "start magic"
            foo( self )
            print "end magic"
        return magic

    @_decorator
    def bar( self ) :
        print "normal call"

test = Test()

test.bar()

This avoids the call to self to access the decorator and leaves it hidden in the class namespace as a regular method.

>>> import stackoverflow
>>> test = stackoverflow.Test()
>>> test.bar()
start magic
normal call
end magic
>>> 

---

edited to answer question in comments:

How to use the hidden decorator in another class

class Test(object):
    def _decorator(foo):
        def magic( self ) :
            print "start magic"
            foo( self )
            print "end magic"
        return magic

    @_decorator
    def bar( self ) :
        print "normal call"

    _decorator = staticmethod( _decorator )

class TestB( Test ):
    @Test._decorator
    def bar( self ):
        print "override bar in"
        super( TestB, self ).bar()
        print "override bar out"

print "Normal:"
test = Test()
test.bar()
print

print "Inherited:"
b = TestB()
b.bar()
print

Output:

Normal:
start magic
normal call
end magic

Inherited:
start magic
override bar in
start magic
normal call
end magic
override bar out
end magic

define a decorator as method inside class

So I was able to draft up something for you, below is the code:

def count(func):
    def wrapper(self):
        TestClass.call_count += 1
        func(self)

    return wrapper

class TestClass(object):
    call_count = 0

    @count
    def hello(self):
        return 'hello'

if __name__ == '__main__':
    x = TestClass()
    for i in range(10):
        x.hello()

    print(TestClass.call_count)

---

Why would it cause problems to have the decorator in a class:

It's not straight forward to have a decorator function inside the class. The reasons are below:

Reason 1
Every class method must take an argument self which is the instance of the class through which the function is being called. Now if you make the decorator function take a self argument, the decorator call @count would fail as it get converted to count() which doesn't pass the self argument and hence the error:

TypeError: wrapper() missing 1 required positional argument: 'self'

Reason 2
Now to avoid that you can make your decorator as static by changing the declaration like below:

@staticmethod
def count(func):
    pass

But then you have another error:

TypeError: 'staticmethod' object is not callable

Which means you can't have a static method as well. If you can't have a static method in a class, you have to pass the self instance to the method but if you pass the self instance to it, the @count decorator call wouldn't pass the self instance and hence it won't work.

So here is a blog that explains it quite well, the issues associated with it and what are the alternatives.

I personally prefer the option to have a helper class to hold all my decorators that can be used instead of the only class in which it's defined. This would give you the flexibility to reuse the decorators instead of redefining them which would follow the ideology

code once, reuse over and over again.

Class decorators for methods in classes

Class decorators accept the function as a subject within the __init__ method (hence the log message), so your decorator code should look like:

class PrintLog(object):

    def __init__(self, function):
        self.function = function

    def __call__(self):
        @wraps(self.function)
        def wrapped(*args):
            print('I am a log')
            return self.function(*args)
        return wrapped

Sorry if this doesn’t work, I’m answering on my mobile device.

EDIT:

Okay so this is probably not what you want, but this is the way to do it:

from functools import update_wrapper, partial, wraps

class PrintLog(object):

    def __init__(self, func):
        update_wrapper(self, func)
        self.func = func
    
    def __get__(self, obj, objtype):
        """Support instance methods."""
        return partial(self.__call__, obj)

    def __call__(self, obj, *args, **kwargs):
        @wraps(self.func)
        def wrapped(*args):
            print('I am a log')
            return self.func(*args)
        return wrapped(obj, *args)

class foo(object):
    def __init__(self, rs: str) -> None:
        self.ter = rs

    @PrintLog
    def baz(self) -> None:
        print('inside baz')

bar = foo('2')
print('running bar.baz()')
bar.baz()

The decorator has to have the __get__ method defined because you're applying the decorator to an instance method. How would a descriptor have the context of the foo instance?

Ref: Decorating Python class methods - how do I pass the instance to the decorator?

Python decorator classes with kwargs move function object

Your use of Decorator1 causes the function being decorated to be passed to Decorator1.__init__ as an argument.

Your use of Decorator2 causes an instance of Decorator2 to be created, and then the function being decorated is passed as an argument to Decorator2.__call__.

Decorator syntax is a shortcut for defining a function, then

1. calling the decorator on that function
2. and binding the result of the decorator call to the original name of the function.

---

That is, the following are very different:

# Decorator1 is called on my_function_1
# Equivalent to my_function_1 = Decorator1(my_function_1)
@Decorator1
def my_function_1(height=2):
    print('This is my_function_1')
    return True

# An instance of Decorator1 is called on my_function_1.
# Equivalent to my_function_1 = Decorator1()(my_function_1)
@Decorator1()
def my_function_1(height=2):
    print('This is my_function_1')
    return True

---

In this sense, decorator syntax differs from class statement syntax, where

class Foo:
    ...

and

class Foo():
    ...

are identical.

Effectively, the string following @ has always been an expression that evaluates to a callable object. Starting in Python 3.9, it can literally be any such expression, rather than being restricted to a dotted name or a dotted-name call as in earlier versions.

Decorating a Python class with a decorator as a class

Won't be a full answer, but I think it's helpful to review the basics of a decorator. This is what decorating looks like:

@Logger
class A:
  # A's code

By definition, it's equivalent to doing this:

class A
  # A's code

A = Logger(A) # Logger has to be callable because...it's called

Sources often say that decorators "modify", but that's really just the intended use. Technically, all you need is A to have a definition (so a function, method, or class) and Logger to be callable. If Logger returned "Hello, World", that's what A becomes.

Okay, let's pretend we didn't decorate A for a bit and think about what it would take for Logger(A) to be "modifying." Well, A is a class, and you call a class to create instances: A(*args). Therefore, Logger(A)(*args) must also be instances of A. But Logger(A) isn't the class A, it's an instance of Logger. Luckily, you can make instances callable by defining the __call__ method in its class. Logger's __call__ method calls the class stored in its cls attribute and returns the instance.

As for parameters in a decorator, it also helps to think about what it's equivalent to. You're interested in doing this:

@Logger(x='y')
class A:
  # A code

So it's equivalent to this:

class A:
  # A code

A = Logger(x = 'y')(A)

Note that Logger itself is not taking A as an argument. It's taking 'y' as an argument and returning another callable that takes A as an argument. So if Logger is a class, Logger(x = 'y') would be a Logger instance. Instances of a class can also serve as decorators if the class has a __call__ method!

Class method decorator with self arguments?

Yes. Instead of passing in the instance attribute at class definition time, check it at runtime:

def check_authorization(f):
    def wrapper(*args):
        print args[0].url
        return f(*args)
    return wrapper

class Client(object):
    def __init__(self, url):
        self.url = url

    @check_authorization
    def get(self):
        print 'get'

>>> Client('http://www.google.com').get()
http://www.google.com
get

The decorator intercepts the method arguments; the first argument is the instance, so it reads the attribute off of that. You can pass in the attribute name as a string to the decorator and use getattr if you don't want to hardcode the attribute name:

def check_authorization(attribute):
    def _check_authorization(f):
        def wrapper(self, *args):
            print getattr(self, attribute)
            return f(self, *args)
        return wrapper
    return _check_authorization

Declaring decorator inside a class

Instead of defining the decorator inside the class you can just intercept the self parameter:

import functools

def enter_exit_info(func):
    @functools.wraps(func)
    def wrapper(self, *arg, **kw):
        print '-- entering', func.__name__
        print '-- ', self.__dict__
        res = func(self, *arg, **kw)
        print '-- exiting', func.__name__
        print '-- ', self.__dict__
        return res
    return wrapper

class TestWrapper():
    def __init__(self, a, b):
        self.a = a
        self.b = b
        self.c = 0
    
    @enter_exit_info
    def add_in_c(self):
        self.c = self.a + self.b
        print self.c

    @enter_exit_info
    def mult_in_c(self):
        self.c = self.a * self.b
        print self.c

if __name__ == '__main__':
    t = TestWrapper(2, 3)
    t.add_in_c()
    t.mult_in_c()

Using decorators with class

From Satwik Kansal’s brilliant Metaprogramming in Python IBM tutorial , I discovered this gem:

Satwik first defined a decorator:

from functools import wraps
import random
import time

def wait_random(min_wait=1, max_wait=30):
    def inner_function(func):
        @wraps(func)
        def wrapper(args, **kwargs):
            time.sleep(random.randint(min_wait, max_wait))
            return func(args, **kwargs)

        return wrapper

    return inner_function

And then he created a class wrapper that will apply this decorator to a class:

def classwrapper(cls):
    for name, val in vars(cls).items():
        #callable return True if the argument is callable
        #i.e. implements the __call
        if callable(val):
            #instead of val, wrap it with our decorator.
            setattr(cls, name, wait_random()(val))
    return cls

Application:

# decorate a function

@wait_random(10, 15)
def function_to_scrape():
    #some scraping stuff

# decorate a class

@classwrapper
class Scraper:
    # some scraping stuff

To make use of it in your case, substitute wait_random decorator with your own. Turn your function to a decorator.
E.g

from functools import wraps
import allure

def apply_allure():
    def inner_function(func):
        @wraps(func)
        def wrapper(args, **kwargs):
            func = allure.step(func)
            return func(args, **kwargs)

        return wrapper

    return inner_function

In the classwrapper replace wait_random with apply_allure:

Do read the tutorial for more information and explanations

---

Related Topics