Creating Decorators with Extra Arguments

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 , rather than making it a 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 of them used method.__name__ and one passed the name to), I have to check inside foo_register to see if the first argument is a decorator, and if so, I should: return decorator(method_name) (instead of return decorator ). This kind of “check to make sure it's called” seems very hacky. Is there a better way to create a multi-user decorator like this?

PS I already know that I can demand that the decorator be called, but this is not a "solution". I want the API to feel natural. My wife loves to decorate, and I do not want to spoil it.

Answer 1

Glenn, I had to do this. I guess I'm glad there is no “magic” way to do this. I hate them.

So here is my own answer (method names are different from above, but the 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 

Usage example (both versions work the same way):

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

Answer 2

The cleanest way I know for 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 a decorator is called without optional arguments like this:

 @decorator def function ... 

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

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

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

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

Python 3

Please note that the decorator's signature above can be enhanced using Python 3 *, specific syntax to ensure that keyword arguments are safely used. Just replace the signature of the outermost function with:

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

Answer 3

By using the answers here and elsewhere, as well as a lot of trial and error, I found that in fact there is a much simpler and more general way to force decorators to accept additional arguments. He checks the arguments with which he was called, but there is no other way to do this.

The key is to decorate your decorator .

Decorator decorator code

Here is the decorator decorator (this code is generic and can be used by anyone who needs an additional 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 

Using

Using it is as simple as:

• Create a decorator as usual.
• After the first argument to the objective function, add optional arguments.
• Decorate the optional_arg_decorator 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, save the attribute of the function with the method name passed in or __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 you can use 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 4

What 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 

Answer 5

Improved general decorator code

Here is my adaptation of @Nicole 's answer with the following improvements:

• optional kwargs can be passed to the decorated decorator
• decorated decorator can be a related 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 

Answer 6

Now, when this old stream rises again upstairs, Lemmy just throw some jewelry:

 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 magic decorator is only in one line!

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

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

Answer 7

A generic decorator for decorating decorator descriptions, expressing that decorated decor accepts default arguments that are set if they are not explicitly specified.

 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 any of the 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): ... 

Answer 8

If you want this function on multiple decorators, you can evade the code template using the decorator for the 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 a drawback when you cannot pass 1 argument as a function to the decorator.

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

Answer 9

Here's another variation that is pretty brief 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 it is possible to call inner_decorator with only one parameter, you can do @decorator , @decorator() , @decorator(24) , etc.

This can be generalized to "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(): ... 

Answer 10

Here is another solution that works as well if the optional argument is 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) 

Thus, both syntaxes 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 

Answer 11

Here is my solution written for python3. It differs from others in that it defines the class to be called, not the 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 need 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 12

I made a simple package to solve the problem

Mounting

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

overview

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

This works for a 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 

A specific 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 related methods

 class O: @hiwrap(hi=False) def fn(self, n): print(n) O().fn(1) #> bye #> 1 

Class methods / properties are also supported.

 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 as call is ok too

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

Answer 13

A similar solution, like checking the type and length of arguments using called 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 

Answer 14

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

It supports two development styles: nested (as in Python decorator factories) and flat (one level less than nesting). Here's how your example will 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 shell is not called at all if the arguments are invalid).

decopatch supports an additional development style, which I call double-flat , which is designed to create shells with a signature-saving function like this. Your example will 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 your code is executed in method calls. This may not be desirable - you may want to do things only once during the decoration - the previous style will be better for this.

You can verify 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.