Making decorators with optional arguments

user_128463 0
Viewed 37776
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.

wrapper decorator python
14 Answers

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.

user_72622470

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)
user_76415700

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
user_52741470

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
user_21624760
Related