Python decorator optional argument

Question

from functools import wraps
def logged(func):
    @wraps(func)
    def with_logging(*args, **kwargs):
        print func.__name__ + " was called"
        return func(*args, **kwargs)
    return with_logging

@logged
def f(x):
   """does some math"""
   return x + x * x

print f.__name__  # prints 'f'
print f.__doc__   # prints 'does some math'

Given this sample code, how would I be able to do @logged(variable) ?

I tried this

from functools import wraps
def logged(func):
    def outer(var):
        @wraps(func)
        def with_logging(*args, **kwargs):
            print func.__name__ + " was called"
            return func(*args, **kwargs)
        return with_logging
    return outer

I was hoping to execute like this: logged(func)(session_variable)

But doesn't work. Any idea? I want to be able to do @logged and @logged(var) ( or even @logged(var1, var2)) Thanks.

Answer 1

The trick here is, you have to introspect what you are given:

def logged(*setting_args, **setting_kwargs):
    no_args = False
    if len(setting_args) == 1 \
        and not setting_kwargs \
        and callable(setting_args[0]):
        # We were called without args
        func = setting_args[0]
        no_args = True

    def outer(func):
        @wraps(func)
        def with_logging(*args, **kwargs):
            print "{} was called".format(func.__name__)
            print "Setting args are: {}".format(setting_args)
            print "Setting keyword args are: {}".format(setting_kwargs)
            return func(*args, **kwargs)
        return with_logging

    if no_args:
        return outer(func)
    else:
        return outer

This will work with any of the following:

# No arguments
@logged
def some_function(x):
    pass

# One or more arguments
@logged(1, 2, 3)
def some_function(x):
    pass

# One or more keyword arguments
@logged(key=1, another_key=2)
def some_function(x):
    pass

# A mix of the two
@logged(1, 2, key=3)
def some_function(x):
    pass

It will not work if it is called with only one callable argument:

# This will break.
@logged(lambda: "Just for fun")
def some_function(x):
    pass

There is no way to tell the difference between a single callable setting and a no-arg invocation of the decorator. However, you can pass a garbage keyword arg to get around even that if you need to:

# This gets around the above limitation
@logged(lambda: "Just for fun", ignored=True)
def some_function(x):
    pass

Answer 2

The question is more than 6 years old, and already has answers. I bumped into same situation - had to update a decorator used at lot of places in the code, and wanted to add an optional argument.

I am able to get it done, using a different approach - from the book Python CookBook 3rd Edition , Chapter 9 - 9.6. Defining a Decorator That Takes an Optional Argument. It takes the problem, presents the solution and end it with a discussion (just great).

---

Solution : for Python 3.3+

from functools import wraps, partial

def logged(func=None, *, var1=None, var2=None):
    if func is None:
        return partial(logged, var1=var1, var2=var2)

    @wraps(func)
    def with_logging(*args, **kwargs):
        print func.__name__ + " was called"
        return func(*args, **kwargs)
    return with_logging

with above you can do either of the following:

@logged
def f(x):

@logger(var1)
def f(x):

@logger(var1, var2)
def f(x) 

---

Explanation (best to look for it in the book)

To understand how the code works, you need to have a firm understanding of how decorators get applied to functions and their calling conventions.

1. Simple decorator such as this:

# Example use
@logged
def add(x, y):
    return x + y

The calling sequence is as follows:

def add(x, y):
    return x + y

add = logged(add)

In this case, the function to be wrapped is simply passed to logged as the first argument. Thus, in the solution, the first argument of logged() is the function being wrapped. All of the other arguments must have default values.

2. Decorator taking arguments:

@logged(level=logging.CRITICAL, name='example')
def spam():
    print('Spam!')

The calling sequence is as follows:

def spam():
    print('Spam!')

spam = logged(level=logging.CRITICAL, name='example')(spam)

The above last line, is how a decorator with arguments is called ie in the initial invocation of logged() the function to be decorated spam() is not passed, thus we set it optional in decorator ie func=None in logged definition. So, in the first call only the arguments are passed.

This, in turn, forces the other arguments to be specified by keyword. Furthermore, when arguments are passed, a decorator is supposed to return a function that accepts the function and wraps it (see Recipe 9.5). To do this, the solution uses a clever trick involving functools.partial . Specifically, it simply returns a partially applied version of itself where all arguments are fixed except for the function to be wrapped .

Answer 3

put the def outer(var) outward, that is

def outer(var):
    def logged(func):
        ...

, then use @outer(somevar) for your function, this would work.

Answer 4

There is minor error in the code you tried. Instead of creating nested functions with arguments as func > var > *args, **kwargs , the order should be var > func > *args, **kwargs .

Below is the code snippet that would suffice your requirement.

from functools import wraps

def logged(var=None):
    def outer(func):
        @wraps(func)
        def with_logging(*args, **kwargs):
            print func.__name__ + " was called"
            return func(*args, **kwargs)
        return with_logging
    return outer

You may call this decorator as:

@logged
def func1():
    ...

OR,

@logged(xyz)
def func2():
    ...

To know more on how decorators work, refer article Decorators with optional arguments .

Answer 5

A possible alternative solution to Sean Vieira's could be:

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 same drawback where you can't pass 1 argument as function to decorator, but if you want this functionality on multiple decorators you can evade the code boilerplate.