创建带有可选参数的装饰器

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

举个例子：下面的代码用 foo_register 装饰了 my_function，而不是直接用 decorator。

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

再举个例子：下面的代码运行效果是正常的。

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

如果我想让它在两个应用中都能正常工作（一个是用 method.__name__，另一个是直接传名字进来），我就得在 foo_register 里面检查第一个参数是不是一个装饰器。如果是的话，我就得这样写：return decorator(method_name)（而不是 return decorator）。这种“检查它是否可以调用”的方式感觉有点笨拙。有没有更好的方法来创建一个多用途的装饰器呢？

附言：我知道我可以要求装饰器必须被调用，但那并不是一个“解决方案”。我希望这个接口用起来更自然。我的妻子喜欢装饰，我不想破坏这个乐趣。

Answer 1

通过这里和其他地方的回答，以及一系列的尝试和错误，我发现其实有一种更简单、更通用的方法来让装饰器接受可选参数。它确实会检查调用时传入的参数，但没有其他的方式来做到这一点。

关键在于给你的装饰器再加一个装饰器。

通用装饰器的装饰器代码

下面是这个装饰器的装饰器（这段代码是通用的，任何需要可选参数装饰器的人都可以使用）：

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

使用方法

使用起来非常简单：

1. 像往常一样创建一个装饰器。
2. 在第一个目标函数参数后面，添加你的可选参数。
3. 用optional_arg_decorator来装饰这个装饰器。

示例：

@optional_arg_decorator
def example_decorator_with_args(fn, optional_arg = 'Default Value'):
    ...
    return fn

测试案例

针对你的使用场景：

所以在你的情况下，要保存一个函数的属性，使用传入的方法名，或者如果是None就用__name__：

@optional_arg_decorator
def register_method(fn, method_name = None):
    fn.gw_method = method_name or fn.__name__
    return fn

添加装饰的方法

现在你有了一个可以带参数或不带参数使用的装饰器：

@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 2

我知道的最简单的方法是这样的：

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

解释

当装饰器没有传入可选参数时，像这样调用：

@decorator
def function ...

这个时候，函数会作为第一个参数传入，装饰器会返回装饰后的函数，这个结果是我们预期的。

如果装饰器传入了一个或多个可选参数，像这样调用：

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

那么装饰器会接收到一个值为 None 的函数参数，因此会返回一个用于装饰的函数，这也是我们预期的结果。

Python 3

需要注意的是，上面的装饰器写法可以用 Python 3 特有的 *, 语法来改进，这样可以确保安全使用关键字参数。只需将最外层函数的签名替换为：

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

Answer 3

Glenn - 我当时必须这样做。我想我很高兴没有什么“魔法”方法可以做到这一点。我讨厌那些。

所以，这是我自己的答案（方法名和上面不同，但概念是一样的）：

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

示例用法（两个版本的效果是一样的）：

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

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