Making decorators with optional arguments [duplicate]

Question

from functools import wraps

def foo_register(method_name=None):
    """Does stuff."""
    def decorator(method):
        if method_name is None:
            method.gw_method = method.__name__
        else:
            method.gw_method = method_name
        @wraps(method)
        def wrapper(*args, **kwargs):
            method(*args, **kwargs)
        return wrapper
    return decorator

Example: The following decorates my_function with foo_register instead of ever making it to decorator.

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

Example: The following works as expected.

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

If I want it to work correctly in both applications (one using method.__name__ and one passing the name in), I have to check inside of foo_register to see if the first argument is a decorator, and if so, I have to: return decorator(method_name) (instead of return decorator). This sort of "check to see if it's a callable" seems very hackish. Is there a nicer way to create a multi-use decorator like this?

P.S. I already know that I can require the decorator to be called, but that's not a "solution". I want the API to feel natural. My wife loves decorating, and I don't want to ruin that.

Roca

That is the answer. The function should either be a decorator function, or a function which returns a decorator function, not magically one or the other depending on its arguments. – Swoon

@Glenn but his wife loves decorating. And this is an interesting challenge. – Manoeuvre

I only had to add two more lines of code (in answer below) which equates to about 180 bytes or so of savings. That means I don't have to buy a new hard-drive, so my wife can keep decorating. – Roca

Answer 1

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...')

Roca

This is a function with radically different behavior depending on its arguments. That's what I mean by magic: it "figures out what you mean" instead of expecting the user to say what he means to begin with. – Swoon

FWIW, the whole concept of decorators is pretty magical. Not like Lucky Charms magical, but magical nonetheless. I think to make the wife really happy, there should be a decorator decorator in this situation that makes a decorator use default arguments if it's invoked with none. Of course this wouldn't work if it's actually passed a callable. – Corena

@Corena - When I'm coding a decorator I feel like I'm watching the van/bridge scene from "Inception". I agree they aren't much fun to maintain, but they sure help to make a library more user friendly (ie, ugly implementation detail). – Roca

@Aaron - why did you delete your answer? It was a great answer (which I just up-voted and attempted to comment on). (comment was: "@Aaron - 1+ clean and explicit FTW.") – Roca

@orokusaki. I decided that I don't like splitting it because it's not an orthogonal split. The register method still knows about names and unless I can think of a better way to do it, then what you have is cleaner. Thanks for the upvote though. – Manoeuvre

@orokusaki: I used to find them really confusing, but getting used to the idea of passing and returning functions, mostly just by doing a lot of it, has helped to make them only kind of confusing. – Corena

Answer 2

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, ...):

Ramiroramjet

If one accepts the requirement of keyword args instead of positional args (which I do), this is by far the best answer. – Kennel

I concur this should be the answer. It's neat and clear. – Newsmonger

You can enforce the keyword args only requirement by adding * to your function definition. E.g., def decorator(original_function=None, *, argument1=None, argument2=None, ...): – Scleroma

@Scleroma I've incorporated your astute Python 3-specific recommendation into the original answer. We could even extend this a bit further with a clever Python 2.7 hack, but... it's hard to see the point. Python 2.7 is on the cusp of its well-deserved deathbed. – Thay

in case decorator() func signature is lost for intellisense, thus programmer will not see func arguments and return value hints when opening bracket ( – Viol

It seems to me that this works on with keyword optional arguments, for example this works: @decorator as does this: @decorator(optional_argument1='xyz') but @decorator('xyz') does not work. Is this correct? Is there a way to make all three work? Thanks. This does seem the cleanest and clearest method I have seen so far. – Sash

Answer 3

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 :)'

Benzoic

Ah yes, figured that would work, was about to write this myself.. +1 – Confine

In case your decorator received a class as a parameter, this solution WON'T work, since callable(args[0]) returns True in both cases. When you're decorating the function AND when you're invoking it. – Tolkan

Answer 4

How about

from functools import wraps, partial

def foo_register(method=None, string=None):
    if not callable(method):
        return partial(foo_register, string=method)
    method.gw_method = string or method.__name__
    @wraps(method)
    def wrapper(*args, **kwargs):
        method(*args, **kwargs)
    return wrapper

Helbona

there is no such thing as dead threads on stackoverflow. if the best answer comes too late for the original asker: tough luck. but for others discovering this later through searching, it’s always valuable to answer if your answer is valuable. – Adeliaadelice

Answer 5

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

Nial

somehow I totally missed this answer when made my own... it's brilliant. – Beam

we should get this into functools -- are you on Python ideas? – Beam

(even though the fn-decorator is not exposed in the end. it's probably worth using "@wraps(fn)" before wrapped_decorator, so that things such as doctests will still run on the decorator) – Beam

I am subscribed to python ideas, but I wonder if this isn't too hacky/fragile for the standard library. For instance, is_bound_method would yield a false positive if the first argument passed to your free function had a property of the same name as the function. – Nial

@MichaelScottCuthbert Thanks for the doctests tip, I was not aware of that issue. – Nial

RyneEverett and @Benzoic Can you provide test case or example of your solution? Thanks – Alena

Answer 6

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().

Chiasma

careful: type is callable. Consider adding and not (type(args[0]) == type and issubclass(args[0], Exception)) to the condition in the case your decorator takes Exceptions as arguments (like this does). – Russell

@Simon: Feel free to edit accordingly :) – Chiasma

Answer 7

A generic decorator for decorating decorator definitions, expressing that decorated decorator accepts default arguments, which are set if none are explicitly given.

from functools import wraps

def default_arguments(*default_args, **default_kwargs):
  def _dwrapper(decorator):
    @wraps(decorator)
    def _fwrapper(*args, **kwargs):
      if callable(args[0]) and len(args) == 1 and not kwargs:
        return decorator(*default_args, **default_kwargs)(args[0])
      return decorator(*args, **kwargs)
    return _fwrapper
  return _dwrapper

It can be used in either of ways.

from functools import lru_cache   # memoization decorator from Python 3

# apply decorator to decorator post definition
lru_cache = (default_arguments(maxsize=100)) (lru_cache)  
# could also be:
#   @default_arguments(maxsize=100)
#   class lru_cache(object):
#     def __init__(self, maxsize):
#       ...
#     def __call__(self, wrapped_function):
#       ...


@lru_cache   # this works
def fibonacci(n):
  ...

@lru_cache(200)   # this also works
def fibonacci(n):
  ...

Rotatory

Answer 8

If you want this functionality on multiple decorators you can evade the code boilerplate with a decorator for a decorator:

from functools import wraps
import inspect


def decorator_defaults(**defined_defaults):
    def decorator(f):
        args_names = inspect.getargspec(f)[0]

        def wrapper(*new_args, **new_kwargs):
            defaults = dict(defined_defaults, **new_kwargs)
            if len(new_args) == 0:
                return f(**defaults)
            elif len(new_args) == 1 and callable(new_args[0]):
                return f(**defaults)(new_args[0])
            else:
                too_many_args = False
                if len(new_args) > len(args_names):
                    too_many_args = True
                else:
                    for i in range(len(new_args)):
                        arg = new_args[i]
                        arg_name = args_names[i]
                        defaults[arg_name] = arg
                if len(defaults) > len(args_names):
                    too_many_args = True
                if not too_many_args:
                    final_defaults = []
                    for name in args_names:
                        final_defaults.append(defaults[name])
                    return f(*final_defaults)
                if too_many_args:
                    raise TypeError("{0}() takes {1} argument(s) "
                                    "but {2} were given".
                                    format(f.__name__,
                                           len(args_names),
                                           len(defaults)))
        return wrapper
    return decorator


@decorator_defaults(start_val="-=[", end_val="]=-")
def my_text_decorator(start_val, end_val):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            return "".join([f.__name__, ' ', start_val,
                            f(*args, **kwargs), end_val])
        return wrapper
    return decorator


@decorator_defaults(end_val="]=-")
def my_text_decorator2(start_val, end_val):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            return "".join([f.__name__, ' ', start_val,
                            f(*args, **kwargs), end_val])
        return wrapper
    return decorator


@my_text_decorator
def func1a(value):
    return value


@my_text_decorator()
def func2a(value):
    return value


@my_text_decorator2("-=[")
def func2b(value):
    return value


@my_text_decorator(end_val=" ...")
def func3a(value):
    return value


@my_text_decorator2("-=[", end_val=" ...")
def func3b(value):
    return value


@my_text_decorator("|> ", " <|")
def func4a(value):
    return value


@my_text_decorator2("|> ", " <|")
def func4b(value):
    return value


@my_text_decorator(end_val=" ...", start_val="|> ")
def func5a(value):
    return value


@my_text_decorator2("|> ", end_val=" ...")
def func5b(value):
    return value


print(func1a('My sample text'))  # func1a -=[My sample text]=-
print(func2a('My sample text'))  # func2a -=[My sample text]=-
print(func2b('My sample text'))  # func2b -=[My sample text]=-
print(func3a('My sample text'))  # func3a -=[My sample text ...
print(func3b('My sample text'))  # func3b -=[My sample text ...
print(func4a('My sample text'))  # func4a |> My sample text <|
print(func4b('My sample text'))  # func4b |> My sample text <|
print(func5a('My sample text'))  # func5a |> My sample text ...
print(func5b('My sample text'))  # func5b |> My sample text ...

Note: it has the drawback where you can't pass 1 argument as function to decorator.

Note2: if you have tips/notes on how to improve this decorator, you could comment at code review: https://codereview.stackexchange.com/questions/78829/python-decorator-for-optional-arguments-decorator

Relaxation

Answer 9

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

Radford

Answer 10

I was incredibly annoyed by this issue and eventually wrote a library to solve it: decopatch.

It supports two development styles: nested (like in python decorator factories) and flat (one less level of nesting). This is how your example would be implemented in flat mode:

from decopatch import function_decorator, DECORATED
from makefun import wraps

@function_decorator
def foo_register(method_name=None, method=DECORATED):
    if method_name is None:
        method.gw_method = method.__name__
    else:
        method.gw_method = method_name

    # create a signature-preserving wrapper
    @wraps(method)
    def wrapper(*args, **kwargs):
        method(*args, **kwargs)

    return wrapper

Note that I use makefun.wraps instead of functools.wraps here so that the signature is fully preserved (the wrapper is not called at all if the arguments are invalid).

decopatch supports an additional development style, that I call double-flat, that is dedicated to creating signature-preserving function wrappers like this one. Your example would be implemented like this:

from decopatch import function_decorator, WRAPPED, F_ARGS, F_KWARGS

@function_decorator
def foo_register(method_name=None,
                 method=WRAPPED, f_args=F_ARGS, f_kwargs=F_KWARGS):
    # this is directly the wrapper
    if method_name is None:
        method.gw_method = method.__name__
    else:
        method.gw_method = method_name

    method(*f_args, **f_kwargs)

Note that in this style, all of your code is executed in calls to method. This might not be desirable - you might wish to perform things once at decoration time only - for this the previous style would be better.

You can check that both styles work:

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

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

Please check the documentation for details.

Rankins

Answer 11

Here's another variation, which is fairly concise and doesn't use functools:

def decorator(*args, **kwargs):
    def inner_decorator(fn, foo=23, bar=42, abc=None):
        '''Always passed <fn>, the function to decorate.
        # Do whatever decorating is required.
        ...
    if len(args)==1 and len(kwargs)==0 and callable(args[0]):
        return inner_decorator(args[0])
    else:
        return lambda fn: inner_decorator(fn, *args, **kwargs)

Depending on whether inner_decorator can be called with only one parameter, one can then do @decorator, @decorator(), @decorator(24) etc.

This can be generalised to a 'decorator decorator':

def make_inner_decorator(inner_decorator):
    def decorator(*args, **kwargs):
        if len(args)==1 and len(kwargs)==0 and callable(args[0]):
            return inner_decorator(args[0])
        else:
            return lambda fn: inner_decorator(fn, *args, **kwargs)
    return decorator

@make_inner_decorator
def my_decorator(fn, a=34, b='foo'):
    ...

@my_decorator
def foo(): ...

@my_decorator()
def foo(): ...

@my_decorator(42)
def foo(): ...

Leclerc

Answer 12

Here is an other solution that work also if the optional argument is a callable:

def test_equal(func=None, optional_value=None):
    if func is not None and optional_value is not None:
        # prevent user to set func parameter manually
        raise ValueError("Don't set 'func' parameter manually")
    if optional_value is None:
        optional_value = 10  # The default value (if needed)

    def inner(function):
        def func_wrapper(*args, **kwargs):
            # do something
            return function(*args, **kwargs) == optional_value

        return func_wrapper

    if not func:
        return inner
    return inner(func)

This way both syntax will work:

@test_equal
def does_return_10():
    return 10

@test_equal(optional_value=20)
def does_return_20():
    return 20

# does_return_10() return True
# does_return_20() return True

Bayly

Answer 13

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)

Jeffery

Answer 14

A similar solution like those checking the type and length of the arguments using callable classes

class decor(object):

def __init__(self, *args, **kwargs):
    self.decor_args = args
    self.decor_kwargs = kwargs

def __call__(self, *call_args, **call_kwargs):

    if callable(self.decor_args[0]) and len(self.decor_args) == 1:
        func = self.decor_args[0]
        return self.__non_param__call__(func, call_args, call_kwargs)
    else:
        func = call_args[0]
        return self.__param__call__(func)


def __non_param__call__(self, func, call_args, call_kwargs):
        print "No args"
        return func(*call_args, **call_kwargs)

def __param__call__(self, func):
    def wrapper(*args, **kwargs):
        print "With Args"
        return func(*args, **kwargs)
    return wrapper



@decor(a)
def test1(a):
    print 'test' + a

@decor
def test2(b):
    print 'test' + b

Tacita