How to Create a Decorator That Can Be Used Either with or Without Parameters

How to create a decorator that can be used either with or without parameters?

I know this question is old, but some of the comments are new, and while all of the viable solutions are essentially the same, most of them aren't very clean or easy to read.

Like thobe's answer says, the only way to handle both cases is to check for both scenarios. The easiest way is simply to check to see if there is a single argument and it is callabe (NOTE: extra checks will be necessary if your decorator only takes 1 argument and it happens to be a callable object):

def decorator(*args, **kwargs):
    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
        # called as @decorator
    else:
        # called as @decorator(*args, **kwargs)

In the first case, you do what any normal decorator does, return a modified or wrapped version of the passed in function.

In the second case, you return a 'new' decorator that somehow uses the information passed in with *args, **kwargs.

This is fine and all, but having to write it out for every decorator you make can be pretty annoying and not as clean. Instead, it would be nice to be able to automagically modify our decorators without having to re-write them... but that's what decorators are for!

Using the following decorator decorator, we can deocrate our decorators so that they can be used with or without arguments:

def doublewrap(f):
    '''
    a decorator decorator, allowing the decorator to be used as:
    @decorator(with, arguments, and=kwargs)
    or
    @decorator
    '''
    @wraps(f)
    def new_dec(*args, **kwargs):
        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
            # actual decorated function
            return f(args[0])
        else:
            # decorator arguments
            return lambda realf: f(realf, *args, **kwargs)

    return new_dec

Now, we can decorate our decorators with @doublewrap, and they will work with and without arguments, with one caveat:

I noted above but should repeat here, the check in this decorator makes an assumption about the arguments that a decorator can receive (namely that it can't receive a single, callable argument). Since we are making it applicable to any generator now, it needs to be kept in mind, or modified if it will be contradicted.

The following demonstrates its use:

def test_doublewrap():
    from util import doublewrap
    from functools import wraps    

    @doublewrap
    def mult(f, factor=2):
        '''multiply a function's return value'''
        @wraps(f)
        def wrap(*args, **kwargs):
            return factor*f(*args,**kwargs)
        return wrap

    # try normal
    @mult
    def f(x, y):
        return x + y

    # try args
    @mult(3)
    def f2(x, y):
        return x*y

    # try kwargs
    @mult(factor=5)
    def f3(x, y):
        return x - y

    assert f(2,3) == 10
    assert f2(2,5) == 30
    assert f3(8,1) == 5*7

Typing decorators that can be used with or without arguments

The issue comes from the first overload (I should have read the pyright message twice!):

@overload
def decorator(arg: F) -> F:
    ...

This overload accepts a keyword parameter named arg, while the implementation does not!

Of course this does not matter in the case of a decorator used with the @decorator notation, but could if it is called like so: fct2 = decorator(arg=fct).

Python >= 3.8

The best way to solve the issue would be to change the first overload so that arg is a positional-only parameter (so cannot be used as a keyword argument):

@overload
def decorator(arg: F, /) -> F:
    ...

With support for Python < 3.8

Since positional-only parameters come with Python 3.8, we cannot change the first overload as desired.

Instead, let's change the implementation to allow for a **kwargs parameter (an other possibility would be to add a keyword arg parameter). But now we need to handle it properly in the code implementation, for example:

def decorator(*args: Any, **kwargs: Any) -> Any:
    if kwargs:
        raise TypeError("Unexpected keyword argument")

    # rest of the implementation here

How to build a decorator with optional parameters?

I found an example, you can use @trace or @trace('msg1','msg2'): nice!

def trace(*args):
    def _trace(func):
        def wrapper(*args, **kwargs):
            print enter_string
            func(*args, **kwargs)
            print exit_string
        return wrapper
    if len(args) == 1 and callable(args[0]):
        # No arguments, this is the decorator
        # Set default values for the arguments
        enter_string = 'entering'
        exit_string = 'exiting'
        return _trace(args[0])
    else:
        # This is just returning the decorator
        enter_string, exit_string = args
        return _trace

Making decorators with optional arguments

Glenn - I had to do it then. I guess I'm glad that there is not a "magic" way to do it. I hate those.

So, here's my own answer (method names different than above, but same concept):

from functools import wraps

def register_gw_method(method_or_name):
    """Cool!"""
    def decorator(method):
        if callable(method_or_name):
            method.gw_method = method.__name__
        else:
            method.gw_method = method_or_name
        @wraps(method)
        def wrapper(*args, **kwargs):
            method(*args, **kwargs)
        return wrapper
    if callable(method_or_name):
        return decorator(method_or_name)
    return decorator

Example usage (both versions work the same):

@register_gw_method
def my_function():
    print('hi...')

@register_gw_method('say_hi')
def my_function():
    print('hi...')

Decorators with parameters?

The syntax for decorators with arguments is a bit different - the decorator with arguments should return a function that will take a function and return another function. So it should really return a normal decorator. A bit confusing, right? What I mean is:

def decorator_factory(argument):
    def decorator(function):
        def wrapper(*args, **kwargs):
            funny_stuff()
            something_with_argument(argument)
            result = function(*args, **kwargs)
            more_funny_stuff()
            return result
        return wrapper
    return decorator

Here you can read more on the subject - it's also possible to implement this using callable objects and that is also explained there.

How to create a Python decorator that can wrap either coroutine or function?

May be you can find better way to do it, but, for example, you can just move your wrapping logic to some context manager to prevent code duplication:

import asyncio
import functools
import time
from contextlib import contextmanager

def duration(func):
    @contextmanager
    def wrapping_logic():
        start_ts = time.time()
        yield
        dur = time.time() - start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if not asyncio.iscoroutinefunction(func):
            with wrapping_logic():
                return func(*args, **kwargs)
        else:
            async def tmp():
                with wrapping_logic():
                    return (await func(*args, **kwargs))
            return tmp()
    return wrapper

Using python decorator with or without parentheses

some_decorator in the first code snippet is a regular decorator:

@some_decorator
def some_method():
    pass

is equivalent to

some_method = some_decorator(some_method)

On the other hand, some_decorator in the second code snippet is a callable that returns a decorator:

@some_decorator()
def some_method():
    pass

is equivalent to

some_method = some_decorator()(some_method)

As pointed out by Duncan in comments, some decorators are designed to work both ways. Here's a pretty basic implementation of such decorator:

def some_decorator(arg=None):
    def decorator(func):
        def wrapper(*a, **ka):
            return func(*a, **ka)
        return wrapper

    if callable(arg):
        return decorator(arg) # return 'wrapper'
    else:
        return decorator # ... or 'decorator'

pytest.fixture is a more complex example.

Use decorators to call function without arguments

In dec2 you are returning the result of calling func1 with the specified arguments, which is not what you want.

What you want is returning a function f which calls func1 with the specified arguments, that is:

def dec1(*args, **kwargs):
    def dec2(func):
        def f():
            return func(*args, **kwargs)
        return f
    return dec2

 

A more detailed explanation:

Remember that the decorator syntax:

@dec1(1, 2, 3)
def func1(val1, val2, val3):
    ...

is syntactically equivalent to:

def func1(val1, val2, val3):
    ...
func1 = dec1(1, 2, 3)(func1)

so the result of dec1(...) (dec2) is called with the decorated function (func1) as argument at the time you decorate the function. So you don't want dec2 to do anything but return a function that will do something later when it's called.

---

Related Topics