第三章 · 函数

装饰器

本节目标

学完这一节，你会知道：

  1. 装饰器解决什么问题
  2. 函数为什么可以作为参数传给另一个函数
  3. @decorator 语法是什么意思
  4. 如何写一个简单的日志装饰器
  5. 为什么装饰器里常见 *args 和 **kwargs

装饰器是进阶内容。第一次学不需要完全吃透，先理解“在不改原函数的情况下，给函数加功能”。

先跑一个例子

新建文件 decorator_demo.py，写入：

def log_call(func):
    def wrapper():
        print("函数开始执行")
        func()
        print("函数执行结束")
    return wrapper


@log_call
def say_hello():
    print("Hello, Python")


say_hello()

运行：

python3 decorator_demo.py

你会看到：

函数开始执行
Hello, Python
函数执行结束

say_hello() 本来只打印一句话，加上装饰器后，前后多了日志。

为什么需要装饰器？

假设很多函数都需要打印“开始”和“结束”。

不使用装饰器时，每个函数都要重复写：

def task_a():
    print("开始")
    print("执行 A")
    print("结束")


def task_b():
    print("开始")
    print("执行 B")
    print("结束")

使用装饰器，可以把重复逻辑抽出来：

@log_call
def task_a():
    print("执行 A")

原函数专心做自己的事，额外功能交给装饰器。

函数也可以当作数据

在 Python 里，函数可以赋值给变量：

def greet():
    print("你好")


hello = greet
hello()

函数也可以作为参数传给另一个函数：

def run(func):
    func()


run(greet)

理解这一点，装饰器就没那么神秘了。

装饰器的基本结构

def my_decorator(func):
    def wrapper():
        print("执行前")
        func()
        print("执行后")
    return wrapper

使用：

@my_decorator
def say_hi():
    print("Hi")

这等价于：

say_hi = my_decorator(say_hi)

wrapper 是包装后的新函数。

装饰带参数的函数

如果原函数有参数，wrapper 也要能接收参数。

def log_call(func):
    def wrapper(*args, **kwargs):
        print(f"调用函数：{func.__name__}")
        result = func(*args, **kwargs)
        print(f"返回结果：{result}")
        return result
    return wrapper


@log_call
def add(a, b):
    return a + b


print(add(3, 5))

*args 和 **kwargs 让装饰器可以适配各种参数形式的函数。

functools.wraps

装饰器会把原函数替换成 wrapper，这可能让函数名和文档丢失。

推荐写法：

from functools import wraps


def log_call(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f"调用：{func.__name__}")
        return func(*args, **kwargs)
    return wrapper

@wraps(func) 会保留原函数的名字和文档信息。

逐行拆解

再看开头的例子：

def log_call(func):

定义一个装饰器，它接收一个函数。

def wrapper():

定义一个新函数，负责包住原函数。

func()

在包装函数里面调用原函数。

return wrapper

把包装后的函数返回。

@log_call

把下面的函数交给 log_call 装饰。

自己改一改

把 decorator_demo.py 改成：

from functools import wraps


def log_call(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f"准备调用 {func.__name__}")
        result = func(*args, **kwargs)
        print(f"{func.__name__} 调用完成")
        return result
    return wrapper


@log_call
def multiply(a, b):
    return a * b


print(multiply(3, 4))

然后继续改：

  1. 给 multiply() 换成除法函数
  2. 在装饰器里打印传入的 args
  3. 再装饰一个 greet(name) 函数

实战：计时器装饰器

新建文件 timer_decorator.py：

import time
from functools import wraps


def timer(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        end = time.perf_counter()
        print(f"{func.__name__} 耗时：{end - start:.4f} 秒")
        return result
    return wrapper


@timer
def count_to(n):
    total = 0
    for i in range(1, n + 1):
        total += i
    return total


print(count_to(1000000))

计时器装饰器适合观察某个函数运行多久。

了解即可：带参数的装饰器

装饰器本身也可以接收参数，但结构会多一层。

from functools import wraps


def repeat(times):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for _ in range(times):
                func(*args, **kwargs)
        return wrapper
    return decorator


@repeat(3)
def say_hi():
    print("Hi")


say_hi()

第一次学到这里，知道它长什么样就够了。

常见应用场景

装饰器常用于：

  • 日志记录
  • 计时
  • 权限检查
  • 缓存
  • 重试
  • Flask 路由，比如 @app.route("/")

以后你看到 @xxx，基本都可以先理解成：给下面的函数加一层功能。

常见错误

1. wrapper 忘记返回结果

错误写法：

def deco(func):
    def wrapper(*args, **kwargs):
        func(*args, **kwargs)
    return wrapper

如果原函数有返回值，这样会丢失结果。

正确写法：

result = func(*args, **kwargs)
return result

2. wrapper 参数写死

def wrapper():
    return func()

这样只能装饰没有参数的函数。更通用的写法是：

def wrapper(*args, **kwargs):
    return func(*args, **kwargs)

3. 忘记导入 wraps

使用 @wraps(func) 前要写：

from functools import wraps

4. 一开始就写太复杂

装饰器很容易越写越绕。先从日志、计时这两类简单装饰器练起。

小练习

练习 1：开始结束装饰器

写一个装饰器，在函数执行前打印“开始”，执行后打印“结束”。

练习 2：参数日志装饰器

写一个装饰器，打印函数收到的 args 和 kwargs。

练习 3：大写返回值

写一个装饰器，如果原函数返回字符串，就把它转成大写。

参考答案

练习 1：

def show_status(func):
    def wrapper(*args, **kwargs):
        print("开始")
        result = func(*args, **kwargs)
        print("结束")
        return result
    return wrapper

练习 2：

def log_args(func):
    def wrapper(*args, **kwargs):
        print(f"args: {args}")
        print(f"kwargs: {kwargs}")
        return func(*args, **kwargs)
    return wrapper

练习 3：

def uppercase_result(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        if isinstance(result, str):
            return result.upper()
        return result
    return wrapper


@uppercase_result
def say():
    return "hello"


print(say())

小结

这一节你学会了：

  1. 装饰器可以在不改原函数的情况下增加功能
  2. 函数可以作为参数传给另一个函数
  3. @decorator 是一种语法糖
  4. wrapper(*args, **kwargs) 是常见通用写法
  5. functools.wraps 可以保留原函数信息

下一章我们会学习数据结构。列表、元组、集合、字典会帮助你组织更多数据。

装饰器像给函数披一件小外套

装饰器第一次见确实容易绕，三层函数套娃也不太讲武德。先抓住核心：不改原函数，也能在前后加功能。把日志和计时器例子跑通就很好，马哥不要求你今天就把套娃拆到分子级。

讨论 (0)

还没有评论，来抢沙发吧！