Python: Implementation of optional argument decorator as class

Question

After reading the excellent Primer on Python Decorators I thought of implementing some of the fancy (advanced) decorators from the article as classes as an exercise.

So for example the decorator with arguments example

def repeat(num_times):
    def decorator_repeat(func):
        @functools.wraps(func)
        def wrapper_repeat(*args, **kwargs):
            for _ in range(num_times):
                value = func(*args, **kwargs)
            return value
        return wrapper_repeat
    return decorator_repeat

could be implemented as a class like this

class Repeat:
    def __init__(self, times):
        self.times = times

    def __call__(self, fn):
        def _wrapper(*args, **kwargs):
            for _ in range(self.times):
                result = fn(*args, **kwargs)
            return result
        return _wrapper

However I seem to be unable to find a class solution for the optional argument decorator example :

def repeat(_func=None, *, num_times=2):
    def decorator_repeat(func):
        @functools.wraps(func)
        def wrapper_repeat(*args, **kwargs):
            for _ in range(num_times):
                value = func(*args, **kwargs)
            return value
        return wrapper_repeat

    if _func is None:
        return decorator_repeat
    else:
        return decorator_repeat(_func)

Is it just me, or is that one rather wicked? XD Would love to see a solution!

Answer 1

You can override the __new__ method to achieve the same behavior:

def __new__(cls, _func=None, *, times=2):
    obj = super().__new__(cls)
    obj.__init__(times)
    if _func is None:
        return obj
    else:
        return obj(_func)

so that both:

@Repeat
def a():
    print('hi')

and:

@Repeat(times=2)
def a():
    print('hi')

output:

hi
hi

Answer 2

Just came across this old question and gave it another try.

I think this is a rather interesting (recursive) solution:

class Repeat:
    
    def __init__(self, fn=None, *, times=2):
        self._fn = fn
        self._times = times

    def _fn_proxy(self, fn):
        self._fn = fn
        return self

    def __call__(self, *args, **kwargs):
        if self._fn:
            for _ in range(self._times):
                result = self._fn(*args, **kwargs)
            return result
        # assertion: if not self._fn, then args[0] must be the decorated function object
        return self._fn_proxy(args[0])

@Repeat
def fun(x,y):
    print(f"{x} and {y} and fun!")

@Repeat(times=4)
def more_fun(x,y):
    print(f"{x} and {y} and even more fun!")

fun(1,2)
print()
more_fun(3,4)