Appearance

Python装饰器

geekbing2023-11-04Pythonpython 装饰器

前言

先提出几个问题:

  • 什么是装饰器?

  • 装饰器有什么作用?

  • 怎么实现装饰器?

什么是装饰器

Python的装饰器可以用来包装目标对象(函数、类、方法)来扩展和修改其行为,同时又不会永久修改可调用对象本身。

绝大多数装饰器是利用函数的闭包原理实现的。

装饰器有什么作用

装饰器的一大用途是将通用的功能应用到现在的类或函数的行为上,这些功能包括:

  • 日志( logging )

  • 访问控制和授权

  • 衡量函数,如执行时间

  • 限制请求速率( rate-limiting )

  • 缓存,等等

Django REST framework中的视图函数装饰器@api_view

Flask中的路由注册装饰器@app.route

FastAPI中的各种路由注册装饰器@app.get@app.post等等

为什么要掌握装饰器

为什么要掌握装饰器,上面提到的内容听起来很抽象,可能很难看出来装饰器在日常工作中能为开发人员来的好处。

下面尝试通过一个实际例子来回答这个问题。

假设在报告生成程序中有 30 个处理业务逻辑的函数。有一天老板走到你的办公桌前说:“我需要你为报告生成器中的每个步骤都添 加输入/输出日志记录的功能,XX公司需要用其来进行审计。我告诉他们我们可以在周三之前完成。”

如果你对Python装饰器掌握的还不错,应该能冷静地应对这个需求,否则你就要血压飙升了。

如果没有装饰器,可能需要花费好几天时间来逐个修改这 30 个函数,在其中添加手动调用日志记录的代码,想想就很悲催~

但如果你了解装饰器,就能很平静且微笑地对老板说:“别担心,我会在今天下午 2 点之前完成。”

然后你开始着手编写一个通用的 @audit_log 装饰器(只有大约 10 行),并将其快速粘贴到每个函数定义的前面。提交代码后就能休息了。

上面的例子可能夸张了一些。但是装饰器确实很强大。对于所有 Python 程序员来说,理解装饰器是一个里程碑。

怎么实现装饰器

概念

简单的装饰器实现是什么样子的呢?简单来说装饰器是可调用的,将可调用对象作为输入并返回另外一个可调用对象。

其实装饰器不限于函数、方法和类,它们也可以用于装饰其他对象,只要被装饰的对象是可以被执行的(也就是可调用的)。

在 Python 中,这通常意味着对象实现了__call__方法。然而,在实际应用中,我们最常见的是装饰器用于函数和方法。

最简单的装饰器

下面这段代码就具有这种特性,可以认为它是最简单的装饰器。

def null_decorator(func):
  """传入一个函数对象,不做任何修改返回
  """
  return func

从中可以看到,null_decorator 是函数,必然是可调用对象。它将另一个可调用对象作为输入,并直接返回。

下面用这个函数包装(或装饰)另一个函数:

def hello():
  return "Hello!"

hello = null_decorator(hello)

输出结果如下:

>>> hello()
Hello!

这个例子中定义了一个 hello 函数,然后立即运行 null_decorator 函数来装饰它。这个例子看起来没什么用,因为 null_decorator 是故意设计的空装饰器。但后面将用这个例子来讲解 Python 中特殊的装饰器语法。

刚刚是在 hello 上显式调用 null_decorator,然后重新分配给 greet 变量,而使用 Python 的@语法糖可以简化这种写法,如下:

@null_decorator
def hello():
  return "Hello!"

## 对于上面的代码,解释器会解释成下面这样的语句
hello = null_decorator(hello)

输出结果如下:

>>> hello()
Hello!

装饰器可以修改行为

熟悉了装饰器语法,咱们来编写一个有实际作用的装饰器来修改被装饰函数的行为。

这个装饰器稍微复杂一些,修改被装饰函数的打印内容:

# 定义装饰器函数
def alter_phrase(func):
    def wrapper():
        print("What's your name?")
        func()
        print("Nice to meet you")
    return wrapper

  
# 装饰原有函数
@alter_phrase
def bing():
    print("My name is Geek Bing")

这个 alter_phrase 装饰器不像之前那样直接返回输入函数,而是在其中定义一个新函数(闭包)。在调用原函数时,新函数会包装原函数来修改其行为。

输出结果如下:

>>> bing()
What's your name?
My name is Geek Bing
Nice to meet you

是不是和你的预期一致,其实就是把一个函数当参数传递到另一个函数,然后再回调。

装饰器通过这种方式来修改可调用对象的行为,无须永久性地修改原对象。可调用对象的行为仅在装饰时才会改变。

利用这种特性可以将可重用的代码块(如日志记录和其他功能)应用于现有的函数和类。

多个装饰器的调用顺序

多个装饰器应用于一个函数,调用顺序从下至上

def add_go(func):
    def wrapper():
        return func() + " Go"
    return wrapper

def add_python(func):
    def wrapper():
        return func() + " Python"
    return wrapper
  
@add_go
@add_python
def hello():
    return "hello"

输出结果如下:

>>> hello()
hello Python Go

从结果中能清楚的看到装饰器应用顺序是从下至上

首先是 @add_python 装饰器包装输入函数,然后 @add_go 装饰器重新包装这个已经装饰过的函数。

如果将上面的例子拆分开来,以传统方式来应用装饰器,那么装饰器函数调用链如下所示:

hello = add_go(add_python(hello))

从中可以看到先应用 add_python 装饰器,然后再由 add_go 装饰器重新包装前一步生成的包装函数。

这意味着堆叠过多的装饰器会对性能产生影响,因为这等同于添加许多嵌套的函数调用。

在一般实践中不是什么问题,但如果在注重性能的代码中经常使用装饰器,那么要注意这一点。

装饰器处理参数的函数

上面的例子都只是包装了简单的无参函数,没有处理输入函数的参数。

之前的装饰器无法应用含有参数的函数。那么如何装饰带有参数的函数呢?

这种情况下,就要使用 Python 中的变长参数*args**kwargs

下面的 proxy 装饰器就用到了这些特性:

def proxy(func):
   def wrapper(*args, **kwargs):
       return func(*args, **kwargs)
    return wrapper
  • 在 wrapper 闭包中使用***操作符收集所有位置参数和关键字参数,并将其存储在变量 args 和 kwargs 中。
  • wrapper 闭包使用*和**参数解包操作符将收集的参数转发到原输入函数。

继续扩展 proxy 装饰器,下面的装饰器在执行时会记录函数参数和结果:

def trace(func):
    """装饰器:用于跟踪函数调用的参数和返回值。

    当被装饰的函数被调用时,此装饰器会打印出函数的参数和返回值。

    :param func: 需要被装饰的函数对象
    :type func: function
    """

    def wrapper(*args, **kwargs):
        """装饰器内部的包装函数,负责打印参数和调用结果。

        :param args: 被装饰函数的位置参数
        :param kwargs: 被装饰函数的关键字参数
        :return: 被装饰函数的原始返回结果
        """
        # 打印被装饰函数的参数列表
        print(f"函数参数: args: {args}, kwargs: {kwargs}")
        # 调用原始函数并获取返回结果
        original_result = func(*args, **kwargs)
        # 打印被装饰函数的名称和它的返回值
        print(f"函数调用: {func.__name__}() 返回了 {original_result!r}")
        # 返回原始函数的返回值
        return original_result

    # 返回包装函数
    return wrapper


@trace
def hello_pro(first_lang: str, second_lang: str) -> str:
    return f"hello {first_lang} {second_lang}"

输出结果如下:

>>> hello_pro("python", second_lang="go")
函数参数: args: ('python',), kwargs: {'second_lang': 'go'}
函数调用: hello_pro() returned 'hello python go'

使用 trace 对函数进行包装后,调用 hello_pro 函数会打印传递给装饰器函数的参数及其返回值。

接受参数的装饰器

接受参数的装饰器会稍微复杂一些,先看下面的代码:

def trace_pro(print_args=False):
    """
    装饰器:可选地打印被装饰函数的参数及其返回值

    :param print_args: 是否打印函数的参数和返回值, 默认为False不打印
    :type print_args: bool, optional
    """

    def decorator(func):
        """
        实际的装饰器函数

        :param func: 被装饰的函数对象
        :type func: function
        """

        def wrapper(*args, **kwargs):
            """
            包装函数,根据 'print_args' 参数决定是否打印函数调用的详细信息

            :param args: 被装饰函数的位置参数
            :param kwargs: 被装饰函数的关键字参数
            :return: 被装饰函数的返回值
            :rtype: 依据被装饰函数的返回值而定
            """
            if print_args:
                # 打印函数的参数列表和关键字参数
                print(f"函数参数: args: {args}, kwargs: {kwargs}")
            # 调用原始函数并获取结果
            original_result = func(*args, **kwargs)
            if print_args:
                # 打印函数的名称和返回值
                print(f"函数调用: {func.__name__}() 返回了 {original_result!r}")
            # 返回原始函数的返回值
            return original_result

        # 返回包装后的函数
        return wrapper

    # 返回装饰器
    return decorator

  
@trace_pro(print_args=True)
def hello_pro(first_lang: str, second_lang: str) -> str:
    return f"hello {first_lang} {second_lang}"

输出结果如下:

>>> hello_pro("python", second_lang="go")
函数参数: args: ('python',), kwargs: {'second_lang': 'go'}
函数调用: hello_pro() returned 'hello python go'

可以看到,为了增加对参数的支持,装饰器在原本两层嵌套函数上又加了一层。

具体来说,下面的装饰器代码:

@trace_pro(print_args=True)
def hello_pro(first_lang: str, second_lang: str) -> str:...

展开后等同于下面的调用:

hello_pro = trace_pro(print_args=True)(hello_pro)

trace_pro(print_args=True) 首先被调用,并返回 decorator 函数。

然后,decorator 函数被调用,并传入原始的 hello_pro 函数作为参数,返回 wrapper 函数。

最终,hello_pro 的名称指向了这个 wrapper 函数,因此当调用 hello_pro 时,实际上是调用了 wrapper 函数。

实现可选参数的装饰器

如果你用嵌套函数来实现装饰器,装饰器接不接受参数,代码有很大区别,带参数的比不带参数的多一层嵌套。

# 接受参数的装饰器
def trace_pro(print_args=False):
    def decorator(func):
        def wrapper(*args, **kwargs):
      ...
        return wrapper
    return decorator
  
# 不接受参数的装饰器
def trace(func):
    def wrapper(*args, **kwargs):
    ...
    return wrapper

当你实现了一个接受参数的装饰器后,即使所有参数都是有默认值的可选参数,你也必须在使用装饰器时加上括号。

@trace_pro(print_args=False)
@trace_pro()

有参数的装饰器提高了它的使用成本,如果使用时忘记添加括号,程序就会报错。

利用仅限关键字参数,可以很方便的做到只使用 @trace_pro 这种写法。

def trace_pro(func=None, *, print_args=False):
    """
    装饰器:可选地打印被装饰函数的参数及其返回值

    :param print_args: 是否打印函数的参数和返回值, 默认为False不打印
    :type print_args: bool, optional
    """

    def decorator(_func):
        """
        实际的装饰器函数

        :param func: 被装饰的函数对象
        :type func: function
        """

        def wrapper(*args, **kwargs):
            """
            包装函数,根据 'print_args' 参数决定是否打印函数调用的详细信息

            :param args: 被装饰函数的位置参数
            :param kwargs: 被装饰函数的关键字参数
            :return: 被装饰函数的返回值
            :rtype: 依据被装饰函数的返回值而定
            """
            if print_args:
                # 打印函数的参数列表和关键字参数
                print(f"函数参数: args: {args}, kwargs: {kwargs}")
            # 调用原始函数并获取结果
            original_result = _func(*args, **kwargs)
            if print_args:
                # 打印函数的名称和返回值
                print(f"函数调用: {_func.__name__}() 返回了 {original_result!r}")
            # 返回原始函数的返回值
            return original_result

        # 返回包装后的函数
        return wrapper
      
    # 如果装饰器带参数使用,此时func为None,返回decorator函数本身。
    if func is None:
        return decorator
    else:
        # 如果装饰器没有带参数直接装饰函数,此时func不为None,需要立即返回wrapper函数。
        return decorator(func)

让我们分步骤理解上述 trace_pro 装饰器的结构:

  1. 当使用装饰器而没有传递任何参数时(即不带括号的情况),例如 @trace_pro,你实际上是将下面的函数传递给了 trace_pro 装饰器。在这种情况下,func 参数是被装饰的函数对象,不是 None。
  2. 当使用装饰器并传递了参数时(即带括号的情况),例如 @trace_pro(print_args=True),你没有立即传递一个函数给 trace_pro。相反,你是在调用 trace_pro 并传递了 print_args 参数。在这种情况下,func 参数默认是 None,因为你还没有提供函数对象给装饰器。

因此,trace_pro 函数体内的条件判断 if func is None: 用于确定装饰器是如何被调用的:

  • 如果 func is None,意味着 trace_pro 被用作带参数的装饰器工厂。这时,它返回 decorator 函数,该函数将会在稍后实际装饰某个函数时被调用。
  • 如果 func is not None,意味着 trace_pro 被用作不带参数的装饰器。这时,它立即应用 decorator 函数到 func 函数上,并返回 decorator(func) 的结果,即 wrapper 函数。

最后的效果是:

  • 如果你写 @trace_pro,那么 trace_pro 函数直接返回 decorator 函数,然后 decorator 函数应用到紧随其后的函数上。
  • 如果你写 @trace_pro(print_args=True),那么 trace_pro 函数返回 decorator 函数,你实际上调用了 decorator 并将紧随其后的函数作为参数传递给它。

其实展开后等同于下面的调用:

# 直接调用不提供任何参数
trace_pro(hello_pro)("python", second_lang="go")

# 提供可选的关键字参数
trace_pro(print_args=True)(hello_pro)("python", second_lang="go")

# 提供括号调用,但不提供任何参数
trace_pro()(hello_pro)("python", second_lang="go")

像上面这样定义装饰器以后,我们就可以通过多种方式来使用了:

# 不提供任何参数
@trace_pro
def hello_pro(first_lang: str, second_lang: str) -> str:...

# 提供可选的关键字参数
@trace_pro(print_args=True)
def hello_pro(first_lang: str, second_lang: str) -> str:...

# 提供括号调用,但不提供任何参数
@trace_pro()
def hello_pro(first_lang: str, second_lang: str) -> str:...

经过上面的例子你会发现,在使用有参数的装饰器时,一共要做两次函数调用,装饰器总共得包含三层嵌套。

正因为如此,有参数装饰器的代码一直难写、难读。不过没关系,接下来会介绍如何用类来实现有参数的装饰器,减少代码的嵌套层级。

用类来实现装饰器

大部分情况下,我们都会选择用函数来实现装饰器,但这并非唯一的方式。

还记得前面提到过的装饰器是可调用的,函数自然是可调用对象,默认类也是调用的,但是类的实例是不可调用的。

class NotCallableClass:
    pass

# 实例化类
not_callable_instance = NotCallableClass()

# 尝试调用实例将会抛出TypeError
try:
    not_callable_instance()
except TypeError as e:
    print(e)  # 输出: 'NotCallableClass' object is not callable

使用__call__方法,实现类的实例也可调用。

class CallableClass:
    def __call__(self, *args, **kwargs):
        print("Instance is called with arguments:", args, kwargs)

# 实例化类
callable_instance = CallableClass()

# 调用实例,不会抛出TypeError
callable_instance(1, 2, key='value')  
# 输出: Instance is called with arguments: (1, 2) {'key': 'value'}

基于类的这个特性,可以用它来实现装饰器。

# 带参数的装饰器
class trace_pro:
    def __init__(self, print_args=False):
        self.print_args = print_args

    def __call__(self, func):
        def decorator(*args, **kwargs):
            if self.print_args:
                # 打印函数的参数列表和关键字参数
                print(f"函数参数: args: {args}, kwargs: {kwargs}")
            # 调用原始函数并获取结果
            original_result = func(*args, **kwargs)
            if self.print_args:
                # 打印函数的名称和返回值
                print(f"函数调用: {func.__name__}() 返回了 {original_result!r}")
            # 返回原始函数的返回值
            return original_result

        return decorator
     
@trace_pro(print_args=True)
def hello_pro(first_lang: str, second_lang: str) -> str:
    return f"hello {first_lang} {second_lang}"

输出结果如下:

>>> hello_pro("python", second_lang="go")
函数参数: args: ('python',), kwargs: {'second_lang': 'go'}
函数调用: hello_pro() returned 'hello python go'

上面这个示例展示了,用类的方式声明一个装饰器。我们可以看到这个类有两个成员:

  1. 一个是__init__(),第一次调用 _deco = trace_pro(print_args=True) 实际是在初始化一个 trace_pro 实例。
  2. 一个是__call__(),第二次调用 hello_pro =_deco(hello_pro) 是在调用 trace_pro 实例,触发__call__方法。

从上面的输出可以看到整个程序的执行顺序,这要比“函数式”的方式更易读一些,里面的嵌套也少了一层。

再来看看不带参数的装饰器:

class trace_pro:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        # 打印函数的参数列表和关键字参数
        print(f"函数参数: args: {args}, kwargs: {kwargs}")
        # 调用原始函数并获取结果
        original_result = self.func(*args, **kwargs)
        # 打印函数的名称和返回值
        print(f"函数调用: {self.func.__name__}() 返回了 {original_result!r}")
        # 返回原始函数的返回值
        return original_result
      
@trace_pro
def hello_pro(first_lang: str, second_lang: str) -> str:
    return f"hello {first_lang} {second_lang}"

代码执行结果:

>>> hello_pro("python", second_lang="go")
函数参数: args: ('python',), kwargs: {'second_lang': 'go'}
函数调用: hello_pro() returned 'hello python go'

展开后等同于下面的调用:

trace_pro(hello_pro)("python", second_lang="go")

类实例化时,被装饰的函数会作为唯一的初始化参数传递到类的__init__方法中。

然后,类的实例调用时触发__call__方法。

至于用类实现可选参数的装饰器,比用函数实现复杂一些,此处就不做讨论了,有兴趣的同学可自行查阅相关资料。

如何编写可调式的装饰器

在使用装饰器时,实际上是用一个函数替换另一个函数,常会出现一些副作用,原函数的元数据会丢失。

包装闭包隐藏了原函数的名称、文档字符串和参数列表:

def alter_phrase(func):
    def wrapper():
        print("What's your name?")
        func()
        print("Nice to meet you")

    return wrapper

def bing():
    """bing函数文档字符串"""
    print("My name is Geek Bing")

decorated_bing = alter_phrase(bing)

输出结果如下:

>>> bing.__name__
'bing'
>>> bing.__doc__
'bing函数文档字符串'
>>> decorated_bing.__name__
'wrapper'
>>> decorated_bing.__doc__
None

上面访问这个函数的任何元数据,看到的都是包装闭包的元数据。

这增加了调试程序和使用 Python 解释器的难度。使用 Python 标准库中的functools.wraps装饰器能避免这个问题。

import functools

def alter_phrase(func):
    @functools.wraps(func)
    def wrapper():
        print("What's your name?")
        func()
        print("Nice to meet you")

    return wrapper

def bing():
    """bing函数文档字符串"""
    print("My name is Geek Bing")

decorated_bing = alter_phrase(bing)

输出结果如下:

>>> bing.__name__
'bing'
>>> bing.__doc__
'bing函数文档字符串'
>>> decorated_bing.__name__
'bing'
>>> decorated_bing.__doc__
'bing函数文档字符串'

functools.wraps 能够将丢失的元数据从被装饰的函数复制到装饰器闭包中。

建议在编写所有装饰器时都使用functools.wraps,这花不了多少时间,同时可以减少自己和其他人的调试难度。

小结

看了上面这么多例子,估计大家有点晕,我们来做个小结吧。

装饰器是 Python 提供的一种糖语法。

表面上看,装饰器就是扩展现有的一个函数的功能,让它可以干一些其他的事,或是在现有的函数功能上再附加上一些别的功能。

往深入了看,我们不难发现,装饰器可以包装所有的可调用对象,任何可调用对象也可以当做装饰器来使用。

装饰器在包装函数的过程中,原始函数的元数据会丢失,你可以通过 functools.wraps 来解决这个问题。

装饰器是一个很有趣且独特的语言特性,可以很容易地将一些非业务功能的、属于控制类型的代码给抽象出来(访问控制和授权,打印日志,函数路由,或是求函数运行时间之类的非业务功能性的代码)。

上次更新 3/25/2024, 10:43:07 AM