Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Using a class to operate both as decorator and decorator factory

Consider the following decorator function, which either returns a decorated function, or a parametrized decorator function:

from functools import wraps, partial, update_wrapper
from inspect import signature

def wrapit(func=None, *, verb='calling'):
    if func is None:  # return a decoratOR
        return partial(wrapit, verb=verb)
    else:  # return a decoratED
        @wraps(func)
        def _func(*args, **kwargs):
            print(f'{verb} {func.__name__} with {args} and {kwargs}')
            return func(*args, **kwargs)
        return _func

Demo:

>>> f = lambda x, y=1: x + y
>>> ff = wrapit(verb='launching')(f)
>>> assert ff(10) == 11
launching <lambda> with (10,) and {}
>>> assert signature(ff) == signature(f)
>>>
>>> # but can also use it as a "decorator factory"
>>> @wrapit(verb='calling')
... def f(x, y=1):
...     return x + y
...
>>> assert ff(10) == 11
launching <lambda> with (10,) and {}
>>> assert signature(ff) == signature(f)

The class form could look something like this:

class Wrapit:
    def __init__(self, func, verb='calling'):
        self.func, self.verb = func, verb
        update_wrapper(self, func)

    def __call__(self, *args, **kwargs):
        print(f'{self.verb} {self.func.__name__} with {args} and {kwargs}')
        return self.func(*args, **kwargs)

But how do we get the class to be able to operate in the "decorator factory" mode that the functional form has (implemented by the if func is None: return partial... How do we integrate that trick in a decorator class?

like image 675
thorwhalen Avatar asked Dec 06 '25 05:12

thorwhalen


1 Answers

As was suggested in the comments, you can do this using the __new__ method:

class Wrapit:
    def __new__(cls, func=None, *, verb='calling'):
        if func is None:
            return partial(cls,verb=verb)
        self = super().__new__(cls)
        self.func, self.verb = func, verb
        update_wrapper(self, func)
        return self

    def __call__(self, *args, **kwargs):
        print(f'{self.verb} {self.func.__name__} with {args} and {kwargs}')
        return self.func(*args, **kwargs)

The __new__ method is called whenever you try to instantiate a class, and the return value of that method is used as the result of the attempted instantiation -- even if it's not an instance of the class!

like image 79
pppery Avatar answered Dec 08 '25 18:12

pppery



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!