装饰器设置函数的属性

Posted 2023-02-23

【中文标题】装饰器设置函数的属性

【英文标题】：decorator to set attributes of function

【发布时间】：2013-03-25 12:52:25

【问题描述】：

我希望仅当登录用户具有所需的权限级别时才能执行不同的功能。

为了让我的生活更简单，我想使用装饰器。下面我尝试在 'decorated' 函数上设置属性 permission - 如下所示。

def permission(permission_required):
    def wrapper(func):
        def inner(*args, **kwargs):
            setattr(func, 'permission_required', permission_required)
            return func(*args, **kwargs)
        return inner
    return wrapper

@permission('user')
def do_x(arg1, arg2):

    ...

@permission('admin')
def do_y(arg1, arg2):
    ...

但是当我这样做时：

fn = do_x
if logged_in_user.access_level == fn.permission_required:
    ...

我收到一个错误'function' object has no attribute 'permission_required'

我错过了什么？

【问题讨论】：

附带说明：我很确定你想在这里使用functools.wraps。不是直接解决您的问题，而是因为当每个函数最终命名为inner、将(*args, **kwargs)、inspect输入错误的源等时，调试这种代码几乎是不可能的。

【参考方案1】：

您正在检查内部（包装）函数的属性，但将其设置在原始（包装）函数上。但是你需要一个包装函数：

def permission(permission_required):
    def decorator(func):
        func.permission_required = permission_required
        return func
    return decorator

你的装饰器需要返回一些东西来替换原来的函数。原始函数本身（添加了属性）就可以了，因为您要做的就是为其添加一个属性。

如果你还需要一个包装器，那么在 包装器函数上设置属性：

from functools import wraps

def permission(permission_required):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # only use a wrapper if you need extra code to be run here
            return func(*args, **kwargs)
        wrapper.permission_required = permission_required
        return wrapper
    return decorator

毕竟，您正在用装饰器返回的包装器替换包装的函数，所以这就是您要在其上查找属性的对象。

我还在包装器中添加了 @functools.wraps() 装饰器，它将重要的识别信息和其他有用的东西从 func 复制到包装器中，使其更易于使用。

【讨论】：

我需要包装器来允许装饰函数中的参数。我使用了你的代码，效果很好 - 但是当我为参数添加包装器时，错误又回来了。

@rikAtee：您不需要包装器来允许装饰函数中的参数。第一个例子只是修改并返回函数；它仍然采用与装饰之前完全相同的参数。

@rikAtee：确实，如果您所做的只是设置属性，那么 不需要 包装器。只有在需要包装器时才添加包装器（例如，添加额外的代码来操作参数或返回值，或者在调用函数时做额外的事情）。

别忘了用@wraps（在functools中）装饰wrapper，这样被装饰的函数就会保留它的属性（包括__doc__）！

【参考方案2】：

你的装饰器应该返回一个可以替换do_x或do_y的函数，而不是返回do_x或do_y的执行结果 你可以修改你的装饰如下：

def permission(permission_required):
    def wrapper(func):
        def inner():
            setattr(func, 'permission_required', permission_required)
            return func
        return inner()
    return wrapper

当然，您还有另一个简单的解决方案：

def permission(permission_required):
    def wrapper(func):
        setattr(func, 'permission_required', permission_required)
        return func
    return wrapper

【讨论】：

你的内部函数什么都不做，所以你用一个函数替换原来的函数，它只是在包装函数上设置一个属性，使它无用.

【参考方案3】：

问题在于，即使您在 inner 中将所需属性设置为包装函数，inner 也会返回装饰函数返回的任何内容，而这通常不是函数本身。

您应该只返回添加了属性的相同原始函数，因此您不必担心这个原始修饰函数可能采用哪些参数，这意味着您可以摆脱其中一个包装级别：

def permission(permission_required):
   def wrapper(func):
       setattr(func, 'permission_required', permission_required)
       return func
   return wrapper

@permission('user')
def do_x(arg1, arg2):
    pass

@permission('admin')
def do_y(arg1, arg2):
    pass

这很好用：

>>> do_x
<function __main__.do_x(arg1, arg2)>
>>> do_x.permission_required
'user'
>>> do_y.permission_required
'admin'