Python 上下文管理器（Context Managers）与 `with` 语句

Python 上下文管理器：让你的代码像管家一样井井有条（with 语句的秘密）

各位观众，各位朋友，欢迎来到“Python魔法学院”！我是今天的魔法讲师，人称“代码界的段子手”—— Dr. Py (虽然我博士论文写的是并行计算，但今天我们不聊秃头话题，聊聊优雅的 Python)。 今天我们要聊一个听起来有点高深，但用起来非常优雅，能让你的代码瞬间变得井井有条的魔法：上下文管理器，以及它背后的得力助手 with 语句。

一、故事的开端：混乱的厨房与优雅的管家

想象一下，你是一个热爱烹饪的美食家，但每次做完饭，厨房都像被龙卷风扫过一样 🌪️。锅碗瓢盆乱七八糟，油烟机上油腻腻，地上洒满了食材碎屑…… 这时候，你是不是特别需要一个能帮你收拾残局的管家？

在 Python 的世界里，上下文管理器就扮演着这个“管家”的角色。它能确保你在执行某些操作前后，自动完成一些必要的准备和清理工作，让你的代码始终保持在一个干净、可控的状态。

比如，打开一个文件，读取数据，然后关闭文件。如果手动操作，你可能会这样写：

file = open("my_data.txt", "r")
try:
    data = file.read()
    # 对数据进行一些操作
    process_data(data)
except Exception as e:
    print(f"发生错误：{e}")
finally:
    file.close()

这段代码看起来没啥问题，但却充满了隐患。如果 file.read() 过程中发生异常，file.close() 可能就无法执行，导致文件资源一直被占用，最终导致资源泄漏。

而有了上下文管理器，我们可以这样写：

with open("my_data.txt", "r") as file:
    data = file.read()
    # 对数据进行一些操作
    process_data(data)

是不是简洁多了？而且，无论 with 语句块中的代码是否发生异常，file.close() 都会被自动调用，确保资源得到释放。 这就是 with 语句的魅力所在，它背后依赖的正是上下文管理器这个强大的工具。

二、什么是上下文管理器？（“魔法”的原理）

简单来说，上下文管理器是一个实现了 __enter__ 和 __exit__ 这两个特殊方法的对象。你可以把它想象成一个“有始有终”的执行单元。

• __enter__(self)： 进入上下文时被调用。这个方法通常用于进行一些准备工作，比如打开文件、获取锁、建立数据库连接等等。它会返回一个值，这个值会被赋值给 with 语句中的 as 后面的变量（比如上面的 file ）。
• __exit__(self, exc_type, exc_val, exc_tb)： 退出上下文时被调用。这个方法负责进行一些清理工作，比如关闭文件、释放锁、断开数据库连接等等。它接收三个参数：
  • exc_type：异常类型，如果没有异常发生，则为 None。
  • exc_val：异常实例，如果没有异常发生，则为 None。
  • exc_tb：异常的回溯信息，如果没有异常发生，则为 None。

__exit__ 方法的返回值决定了是否要抑制异常的传播。如果返回 True，则表示异常已经被处理，不会再向上层抛出；如果返回 False (或者 None)，则表示异常需要继续传播。

让我们用一个表格来总结一下：

方法	作用	何时调用
__enter__	进入上下文，进行准备工作，返回值给 as	在 with 语句块开始执行之前。
__exit__	退出上下文，进行清理工作，处理异常	在 with 语句块执行完毕之后 (无论是否发生异常)。 如果 with 语句块中发生异常，则 exc_type、exc_val、exc_tb 会被赋值；否则，它们都为 None。 __exit__ 的返回值决定了是否抑制异常的传播。 True 抑制， False (或 None) 不抑制。

三、自己动手，打造一个上下文管理器（DIY乐趣）

理论讲完了，让我们来动手实现一个简单的上下文管理器。 比如，我们要创建一个计时器，用来测量代码块的执行时间。

import time

class Timer:
    def __enter__(self):
        self.start_time = time.time()
        return self # 返回自身，可以赋值给 with 语句中的 as 变量

    def __exit__(self, exc_type, exc_val, exc_tb):
        end_time = time.time()
        execution_time = end_time - self.start_time
        print(f"代码块执行时间：{execution_time:.4f} 秒")
        return False # 不抑制异常，让异常继续传播

# 使用 Timer 上下文管理器
with Timer():
    # 模拟一些耗时操作
    time.sleep(1)
    print("耗时操作完成！")

运行这段代码，你会看到类似这样的输出：

耗时操作完成！
代码块执行时间：1.0012 秒

这个 Timer 类就是一个简单的上下文管理器。__enter__ 方法记录了开始时间，__exit__ 方法计算了执行时间并输出。

我们再来看一个稍微复杂一点的例子，这次我们模拟一个数据库连接的上下文管理器。

class DatabaseConnection:
    def __init__(self, database_name):
        self.database_name = database_name
        self.connection = None

    def __enter__(self):
        try:
            self.connection = connect_to_database(self.database_name)  # 假设有这样一个函数
            print(f"成功连接到数据库：{self.database_name}")
            return self.connection
        except Exception as e:
            print(f"连接数据库失败：{e}")
            raise # 重新抛出异常，避免程序继续执行

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.connection:
            try:
                self.connection.close()
                print(f"成功关闭数据库连接：{self.database_name}")
            except Exception as e:
                print(f"关闭数据库连接失败：{e}")
        return False # 不抑制异常

# 使用 DatabaseConnection 上下文管理器
with DatabaseConnection("my_database") as conn:
    # 在这里使用数据库连接进行操作
    if conn:
        try:
            # 执行一些数据库操作，比如查询、更新等等
            print("正在执行数据库操作...")
            # result = conn.execute("SELECT * FROM users")
            # ...
        except Exception as e:
            print(f"数据库操作失败：{e}")

在这个例子中，__enter__ 方法尝试连接到数据库，如果连接成功，则返回连接对象；如果连接失败，则抛出异常。__exit__ 方法负责关闭数据库连接，无论是否发生异常。

注意： 在 __enter__ 中，如果连接失败，一定要重新抛出异常，否则程序会继续执行 with 语句块中的代码，这可能会导致不可预测的结果。

四、 contextlib 模块：偷懒的艺术（内置神器）

如果你不想自己手动实现 __enter__ 和 __exit__ 方法，Python 的 contextlib 模块提供了一些更方便的工具。 它就像一个魔法工具箱，里面装满了各种“上下文管理器生成器”。

1. @contextmanager 装饰器：化繁为简

@contextmanager 装饰器可以将一个生成器函数变成一个上下文管理器。生成器函数需要使用 yield 语句来分隔准备工作和清理工作。

让我们用 @contextmanager 来改造一下之前的 Timer 例子：

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start_time = time.time()
    yield # yield 语句分隔准备工作和清理工作
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"代码块执行时间：{execution_time:.4f} 秒")

# 使用 timer 上下文管理器
with timer():
    # 模拟一些耗时操作
    time.sleep(1)
    print("耗时操作完成！")

是不是更加简洁了？ @contextmanager 装饰器会自动帮你创建 __enter__ 和 __exit__ 方法。 yield 语句之前的代码会在 __enter__ 方法中执行，yield 语句之后的代码会在 __exit__ 方法中执行。

2. suppress：优雅地忽略异常

有时候，我们希望忽略某些特定的异常，让程序继续执行。 contextlib.suppress 可以帮助我们实现这个功能。

from contextlib import suppress

with suppress(FileNotFoundError):
    # 尝试删除一个文件，如果文件不存在，则忽略异常
    os.remove("non_existent_file.txt")
    print("文件删除成功！") # 如果文件存在，才会执行这行代码

在这个例子中，如果 os.remove 抛出 FileNotFoundError 异常，suppress 上下文管理器会捕获这个异常，并阻止它向上层传播。

3. redirect_stdout 和 redirect_stderr：重定向输出

contextlib.redirect_stdout 和 contextlib.redirect_stderr 可以将标准输出和标准错误重定向到其他文件或对象。

import sys
from contextlib import redirect_stdout, redirect_stderr

# 将标准输出重定向到一个文件
with open("output.txt", "w") as f:
    with redirect_stdout(f):
        print("这段文字会被写入到 output.txt 文件中")

# 将标准错误重定向到标准输出
with redirect_stderr(sys.stdout):
    print("这段文字会被当作错误信息输出到标准输出")

这两个上下文管理器在调试和测试时非常有用。

五、 上下文管理器的应用场景（用武之地）

上下文管理器在 Python 中有着广泛的应用，以下是一些常见的场景：

• 文件操作： 确保文件在使用完毕后被正确关闭，防止资源泄漏。
• 锁机制： 获取和释放锁，保证线程安全。
• 数据库连接： 建立和断开数据库连接，管理数据库事务。
• 网络连接： 建立和断开网络连接，处理网络请求。
• 性能测试： 测量代码块的执行时间。
• 临时更改全局状态： 比如临时修改环境变量、全局配置等。

总而言之，任何需要在执行某些操作前后进行准备和清理工作的场景，都可以使用上下文管理器来简化代码，提高代码的可读性和可维护性。

六、 with 语句的进阶用法（更上一层楼）

with 语句可以嵌套使用，也可以同时管理多个上下文管理器。

with open("file1.txt", "r") as f1, open("file2.txt", "w") as f2:
    # 同时读取 file1.txt 和写入 file2.txt
    data = f1.read()
    f2.write(data)

在这个例子中，我们同时打开了两个文件，并使用 with 语句确保它们在使用完毕后被正确关闭。

七、 总结：优雅的代码，从上下文管理器开始

各位观众，今天的“Python魔法学院”就到这里了。希望通过今天的学习，大家能够掌握上下文管理器这个强大的工具，让你的代码像管家一样井井有条，优雅而高效。 记住，好的代码不仅要能运行，更要易于理解和维护。 而上下文管理器，正是提升代码质量的一大利器。

最后，送给大家一句名言：

"Write code that is easy to delete, not easy to extend." – David Heinemeier Hansson

希望大家都能写出简洁、优雅、易于删除的代码！ 谢谢大家！ 👏