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