Making decorators with optional arguments [duplicate]

Answer 1

I've made a simple package to solve the problem

Installation

Master branch

pip install git+https://github.com/ferrine/biwrap

Latest release

pip install biwrap

Overview

Some wrappers may have optional arguments and we often want to avoid @wrapper() calls and use @wrapper instead.

This works for simple wrapper

import biwrap

@biwrap.biwrap
def hiwrap(fn, hi=True):
    def new(*args, **kwargs):
        if hi:
            print('hi')
        else:
            print('bye')
        return fn(*args, **kwargs)
    return new

Defined wrapper can be used in both ways

@hiwrap
def fn(n):
    print(n)
fn(1)
#> hi
#> 1

@hiwrap(hi=False)
def fn(n):
    print(n)
fn(1)
#> bye
#> 1

biwrap also works for bound methods

class O:
    @hiwrap(hi=False)
    def fn(self, n):
        print(n)

O().fn(1)
#> bye
#> 1

Class methods / properties are supported too

class O:
    def __init__(self, n):
        self.n = n

    @classmethod
    @hiwrap
    def fn(cls, n):
        print(n)

    @property
    @hiwrap(hi=False)
    def num(self):
        return self.n
    
    
o = O(2)
o.fn(1)
#> hi
#> 1
print(o.num)
#> bye
#> 2

Function like call is OK too

def fn(n):
    print(n)

fn = hiwrap(fn, hi=False)
fn(1)
#> bye
#> 1

Answer 2

Through the help of the answers here and elsewhere and a bunch of trial and error I've found that there is actually a far easier and generic way to make decorators take optional arguments. It does check the args it was called with but there isn't any other way to do it.

The key is to decorate your decorator.

Generic decorator decorator code

Here is the decorator decorator (this code is generic and can be used by anyone who needs an optional arg decorator):

def optional_arg_decorator(fn):
    def wrapped_decorator(*args):
        if len(args) == 1 and callable(args[0]):
            return fn(args[0])

        else:
            def real_decorator(decoratee):
                return fn(decoratee, *args)

            return real_decorator

    return wrapped_decorator

Usage

Using it is as easy as:

1. Create a decorator like normal.
2. After the first target function argument, add your optional arguments.
3. Decorate the decorator with optional_arg_decorator

Example:

@optional_arg_decorator
def example_decorator_with_args(fn, optional_arg = 'Default Value'):
    ...
    return fn

Test cases

For your use case:

So for your case, to save an attribute on the function with the passed-in method name or the __name__ if None:

@optional_arg_decorator
def register_method(fn, method_name = None):
    fn.gw_method = method_name or fn.__name__
    return fn

Add decorated methods

Now you have a decorator that is usable with or without args:

@register_method('Custom Name')
def custom_name():
    pass

@register_method
def default_name():
    pass

assert custom_name.gw_method == 'Custom Name'
assert default_name.gw_method == 'default_name'

print 'Test passes :)'

Answer 3

Now that this old thread is back at the top anyway, lemme just throw in some Decorator-ception:

def magical_decorator(decorator):
    @wraps(decorator)
    def inner(*args, **kw):
        if len(args) == 1 and not kw and callable(args[0]):
            return decorator()(args[0])
        else:
            return decorator(*args, **kw)
    return inner

Now your magical decorator is just a single line away!

@magical_decorator
def foo_register(...):
    # bla bla

By the way, this works for any decorator. It just causes @foo to behave (as close as possibly) like @foo().

Answer 4

Here's my solution, written for python3. It has a different approach from the other ones since it defines a callable class rather than a function.

class flexible_decorator:

    def __init__(self, arg="This is default"):
        self.arg = arg

    def __call__(self, func):

        def wrapper(*args, **kwargs):
            print("Calling decorated function. arg '%s'" % self.arg)
            func(*args, **kwargs)

        return wrapper

You still have to explicitly call the decorator

@flexible_decorator()
def f(foo):
    print(foo)


@flexible_decorator(arg="This is not default")
def g(bar):
    print(bar)

Answer 5

The cleanest way I know of for doing this is the following:

import functools


def decorator(original_function=None, optional_argument1=None, optional_argument2=None, ...):

    def _decorate(function):

        @functools.wraps(function)
        def wrapped_function(*args, **kwargs):
            ...

        return wrapped_function

    if original_function:
        return _decorate(original_function)

    return _decorate

Explanation

When the decorator is called with no optional arguments like this:

@decorator
def function ...

The function is passed as the first argument and decorate returns the decorated function, as expected.

If the decorator is called with one or more optional arguments like this:

@decorator(optional_argument1='some value')
def function ....

Then decorator is called with the function argument with value None, so a function that decorates is returned, as expected.

Python 3

Note that the decorator signature above may be improved with Python 3-specific *, syntax to enforce safe use of keyword arguments. Simply replace the signature of the outermost function with:

def decorator(original_function=None, *, optional_argument1=None, optional_argument2=None, ...):

Answer 6

Enhanced Generic Decorator Decorator Code

Here's my adaption of @Nicole's answer with the following enhancements:

• optional kwargs may be passed to the decorated decorator
• the decorated decorator may be a bound method

import functools

def optional_arg_decorator(fn):
    @functools.wraps(fn)
    def wrapped_decorator(*args, **kwargs):
        is_bound_method = hasattr(args[0], fn.__name__) if args else False

        if is_bound_method:
            klass = args[0]
            args = args[1:]

        # If no arguments were passed...
        if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
            if is_bound_method:
                return fn(klass, args[0])
            else:
                return fn(args[0])

        else:
            def real_decorator(decoratee):
                if is_bound_method:
                    return fn(klass, decoratee, *args, **kwargs)
                else:
                    return fn(decoratee, *args, **kwargs)
            return real_decorator
    return wrapped_decorator