使用装饰器管理函数内部的属性

Question

这是我现在正在工作的更大代码的玩具示例。假设我有一个函数，并且我想以一种可以记录有效、遗留或禁止的参数的方式来装饰它：

@option("1")
@option("2", legacy=True)
def do_stuff(opt):
    print(f"You chose {opt}!")


do_stuff("1")
do_stuff("2")
do_stuff("3")

这样的结果会是这样的

You chose 1!
2 is deprecated!
You chose 2!
3 is not allowed!

我现在尝试用一些不同的方式来实现这种行为，我最近的尝试是这样的：

from functools import wraps


def option(opt, legacy=False):
    def add_opt(f, option):
        if not hasattr(f, "opts"):
            f.opts = []
        f.opts.append(option)

    def add_legacy_opt(f, option):
        if not hasattr(f, "legacy"):
            f.legacy = []
        f.legacy.append(option)

    def decorate(func):
        @wraps(func)
        def wrapper(option):
            if hasattr(wrapper, "opts") and option not in wrapper.opts:
                print(f"{option} is not allowed!")
                return
            if hasattr(wrapper, "legacy") and option in wrapper.legacy:
                print(f"{option} is deprecated!")
            return func(option)

        add_opt(wrapper, opt)
        if legacy:
            add_legacy_opt(wrapper, opt)

        return wrapper

    return decorate

它几乎有效。问题是它给了我双重遗留信息：

You chose 1!
2 is deprecated!
2 is deprecated!
You chose 2!
3 is not allowed!

如果有人有任何想法，请告诉我！谢谢：D

更新：以不同的顺序进行装饰：

@option("2", legacy=True)
@option("1")
def do_stuff(opt):
    print(f"You chose {opt}!")

Answer 1

您已经围绕包装器创建了一个包装器。 wraps 装饰器实际上使用“更深层次”的函数属性更新每个包装器，这在这里实际上很关键，但是由于每个包装器都保留自己的 .opts 和 .legacy （指向相同的列表对象）所以每一层包装器将 print(f"{option} is deprecated!") 然后 return func(option) 这只是调用嵌套包装器，它会通过相同的检查。

所以观察：

In [1]: from functools import wraps
   ...:
   ...:
   ...: def option(opt, legacy=False):
   ...:     def add_opt(f, option):
   ...:         if not hasattr(f, "opts"):
   ...:             f.opts = []
   ...:         f.opts.append(option)
   ...:
   ...:     def add_legacy_opt(f, option):
   ...:         if not hasattr(f, "legacy"):
   ...:             f.legacy = []
   ...:         f.legacy.append(option)
   ...:
   ...:     def decorate(func):
   ...:         @wraps(func)
   ...:         def wrapper(option):
   ...:             if hasattr(wrapper, "opts") and option not in wrapper.opts:
   ...:                 print(f"{option} is not allowed!")
   ...:                 return
   ...:             if hasattr(wrapper, "legacy") and option in wrapper.legacy:
   ...:                 print(f"{option} is deprecated!")
   ...:             return func(option)
   ...:
   ...:         add_opt(wrapper, opt)
   ...:         if legacy:
   ...:             add_legacy_opt(wrapper, opt)
   ...:
   ...:         return wrapper
   ...:
   ...:     return decorate
   ...:

In [2]: @option("2", legacy=True)
   ...: @option("1")
   ...: def do_stuff(opt):
   ...:     print(f"You chose {opt}!")
   ...:

In [3]: vars(do_stuff)
Out[3]:
{'__wrapped__': <function __main__.do_stuff(opt)>,
 'opts': ['1', '2'],
 'legacy': ['2']}

In [4]: vars(do_stuff.__wrapped__)
Out[4]: {'__wrapped__': <function __main__.do_stuff(opt)>, 'opts': ['1', '2']}

在hasattr(wrapper, "opts") and option not in wrapper.opts 为真的情况下，它只是return 而不调用func。 （抛开设计不谈，这几乎肯定只会引发值错误）。

以相反的顺序装饰“有效”，因为您最后添加了legacy，因此只有最外层的包装器具有.legacy 属性。

一种解决方案是以某种方式识别已被修饰的函数，它们已经被包装，而不是创建另一个包装器，只需更新包装器上的必要属性并返回原始包装器。

首先，让我们将add_opt 和add_legacy 移到外部，以消除装饰器内部的混乱，没有充分的理由将它们定义在内部：

from functools import wraps

def _add_opt(f, option):
    if not hasattr(f, "opts"):
        f.opts = []
    f.opts.append(option)

def _add_legacy_opt(f, option):
    if not hasattr(f, "legacy"):
        f.legacy = []
    f.legacy.append(option)

然后我们可以做如下的事情：

def option(opt, legacy=False):

    def decorate(func):
        if hasattr(func, 'opt') or hasattr(func, 'legacy'): # has been decorated
            wrapper = func # keep old wrapper
        else: # wrap for the first time
            @wraps(func)
            def wrapper(option):
                if hasattr(wrapper, "opts") and option not in wrapper.opts:
                    print(f"{option} is not allowed!")
                    return
                if hasattr(wrapper, "legacy") and option in wrapper.legacy:
                    print(f"{option} is deprecated!")
                return func(option)

        _add_opt(wrapper, opt)
        if legacy:
            _add_legacy_opt(wrapper, opt)

        return wrapper

    return decorate

老实说，这种依赖函数属性的方法看起来很脆弱。但也许你可以想出一个更强大的方法，但这至少应该澄清问题并给你一个解决方案。