CODEWITHSUNDEEP
pythondecorator
David Wolever
David Wolever

Reputation: 154564

Python: shortcut for writing decorators which accept arguments?

Does the Python standard library have a shortcut for writing decorators which accept arguments?

For example, if I want to write a decorator like with_timeout(timeout):

@with_timeout(10.0)
def cook_eggs(eggs):
    while not eggs.are_done():
        eggs.cook()

I have to write something like:

def with_timeout(timeout):
    _func = [None]
    def with_timeout_helper(*args, **kwargs):
        with Timeout(timeout):
            return _func[0](*args, **kwargs)
    def with_timeout_return(f):
        return functools.wraps(f)(with_timeout_helper)
    return with_timeout_return

But that's awfully verbose. Is there a shortcut which makes decorators which accept arguments easier to write?

Note: I realize that it's also possible to use three nested functions to implement decorators with arguments… But that feels just a bit suboptimal too.

For example, possibly something like a @decorator_with_arguments function:

@decorator_with_arguments
def timeout(f, timeout):
    @functools.wraps(f)
    def timeout_helper(*args, **kwargs):
        with Timeout(timeout):
            return f(*args, **kwargs)
    return timeout_helper

Upvotes: 9

Views: 1865

Answers (5)

user1406431
user1406431

Reputation: 1

Another take, without using lambdas:

def decorator_with_arguments(f):
    @functools.wraps(f)
    def with_arguments_helper(*args, **kwargs):
        def decorator(g):
            return f(g, *args, **kwargs)
        return decorator
    return with_arguments_helper

Upvotes: 0

David Wolever
David Wolever

Reputation: 154564

Based on Jakob's suggestion, I've implemented a small Decorator class, which I feel does a fairly decent job:

class Decorator(object):
    def __call__(self, f):
        self.f = f
        return functools.wraps(f)(lambda *a, **kw: self.wrap(*a, **kw))

    def wrap(self, *args, **kwrags):
        raise NotImplemented("Subclasses of Decorator must implement 'wrap'")

class with_timeout(Decorator):
    def __init__(self, timeout):
        self.timeout = timeout

    def wrap(self, *args, **kwargs):
        with Timeout(timeout):
            return self.f(*args, **kwargs)

Upvotes: 4

Karl Knechtel
Karl Knechtel

Reputation: 61607

First, we can define a little meta-decorator:

def decorator_with_arguments(wrapper):
    return lambda *args, **kwargs: lambda func: wrapper(func, *args, **kwargs)

That allows us to create decorators that accept arguments like so:

@decorator_with_arguments
def my_wrapper(func, *decorator_args, **decorator_kwargs):
    def wrapped(*call_args, **call_kwargs):
        print "from decorator:", decorator_args, decorator_kwargs
        func(*call_args, **call_kwargs)
    return wrapped

Which can then be used normally:

@my_wrapper(1, 2, 3)
def test(*args, **kwargs):
    print "passed directly:", args, kwargs

test(4, 5, 6)

Adding functools.wraps decoration is left as an exercise :)

Upvotes: 0

Nathan
Nathan

Reputation: 2995

I know you said it feels suboptimal but I still feel that using three nested models is the cleanest solution. The inner two functions are just the 'normal' way of defining a decorator for a function that takes arguments (see example in python's docs for @wraps). The outer one is really just a function that takes and argument and returns a decorator.

def with_timeout(timeout):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            with Timeout(timeout):
                return f(*args, **kwargs)
        return wrapper
    return decorator

Upvotes: 4

Jakob Bowyer
Jakob Bowyer

Reputation: 34718

I tend to write my decorators as classes to be honest

class TestWithArgs(object):
    def __init__(self, *deco_args, **deco_kwargs):
        self.deco_args = deco_args
        self.deco_kwargs = deco_kwargs
    def __call__(self, func):
        def _wrap(self, *args, **kwargs):
            print "Blah blah blah"
            return func(*args, **kwargs)
        return _wrap

Its nothing if not slightly clearer

Upvotes: 7

Related Questions