Python高级技术之：`Python`的`context manager`：`__enter__`和`__exit__`方法在资源管理中的作用。

各位观众老爷，大家好！今天给大家伙儿聊聊Python里一个非常有意思，但又经常被忽略的小秘密——Context Manager，也就是咱们常说的“上下文管理器”。这玩意儿，说白了，就是来帮你优雅地搞定资源管理的，让你的代码更干净、更安全，也更Pythonic。

啥是资源管理？ 为啥要优雅地搞？

咱们先来说说资源管理。在编程的世界里，资源可不是指你在游戏里攒的金币。 这里的资源，指的是那些用了就得还回去的东西，比如：

• 文件: 打开了总得关掉吧？
• 网络连接: 连接建立了，不用了总得断开吧？
• 数据库连接: 连上了数据库，操作完了总得释放连接吧？
• 锁: 拿到了锁，用完了总得释放吧？

要是你打开了文件，用完了忘了关，或者拿到了锁，用完了忘了放，那可就惨了。轻则程序运行效率下降，重则系统崩溃、数据丢失，甚至被黑客攻击，后果不堪设想。

所以，资源管理非常重要。那为啥要优雅地搞呢？因为直接用try...finally虽然能保证资源被释放，但是代码看起来比较冗余，尤其是嵌套使用的时候，简直就是噩梦。

Context Manager闪亮登场！

Context Manager就是来拯救你的。它可以让你用一种更简洁、更优雅的方式来管理资源，保证资源在使用完毕后一定会被释放，即使在发生异常的情况下也不例外。

Context Manager的秘密武器：__enter__和__exit__

Context Manager的核心在于两个特殊方法：__enter__ 和 __exit__。 它们就像是资源的“开门大吉”和“送客还家”的管家。

• __enter__(self): 这个方法会在进入 with 语句块的时候被调用。它的主要任务是获取资源，并返回一个值（这个值可以是你自己，也可以是其他任何对象，甚至可以是 None）。这个返回值会被赋值给 with 语句中的 as 变量（如果存在的话）。

• __exit__(self, exc_type, exc_val, exc_tb): 这个方法会在退出 with 语句块的时候被调用。它的主要任务是释放资源。无论 with 语句块是正常执行完毕，还是因为发生了异常而退出，__exit__ 都会被执行。

  • exc_type：异常类型。如果没有异常发生，它的值是 None。
  • exc_val：异常实例。如果没有异常发生，它的值是 None。
  • exc_tb：异常追踪对象。如果没有异常发生，它的值是 None。

  如果 __exit__ 方法返回 True，则表示该异常已经被处理，程序会继续执行；如果返回 False 或者 None，则表示该异常没有被处理，程序会抛出异常。

用代码说话：自定义一个Context Manager

光说不练假把式，咱们来写一个简单的Context Manager，模拟一个文件操作：

class FileManager:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.file = None  # 初始化文件对象

    def __enter__(self):
        print("准备打开文件...")
        self.file = open(self.filename, self.mode)
        return self.file  # 返回文件对象，with语句会把它赋值给变量

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("准备关闭文件...")
        if self.file:
            self.file.close()
        if exc_type:
            print(f"发生异常: {exc_type}, {exc_val}")
            # 可以选择处理异常，这里简单地打印出来
        print("文件已关闭")
        return False # 返回False，表示不阻止异常抛出

现在，我们可以这样使用这个Context Manager：

with FileManager("test.txt", "w") as f:
    print("开始写入数据...")
    f.write("Hello, Context Manager!")
    # 故意制造一个异常
    #raise ValueError("这是一个测试异常")
    print("数据写入完成")

print("程序执行完毕")

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

准备打开文件...
开始写入数据...
数据写入完成
准备关闭文件...
文件已关闭
程序执行完毕

如果把 #raise ValueError("这是一个测试异常") 的注释去掉，你会看到：

准备打开文件...
开始写入数据...
准备关闭文件...
发生异常: <class 'ValueError'>, 这是一个测试异常
文件已关闭
Traceback (most recent call last):
  File "...", line ..., in <module>
    raise ValueError("这是一个测试异常")
ValueError: 这是一个测试异常

可以看到，即使发生了异常，__exit__ 方法仍然会被执行，文件会被关闭。

再来一个例子：模拟数据库连接

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

    def __enter__(self):
        print("建立数据库连接...")
        # 这里模拟建立数据库连接，实际情况需要使用相应的数据库驱动
        self.connection = "模拟的数据库连接"
        return self.connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("关闭数据库连接...")
        # 这里模拟关闭数据库连接
        if self.connection:
            self.connection = None
        if exc_type:
            print(f"数据库操作发生异常: {exc_type}, {exc_val}")
        print("数据库连接已关闭")
        return False

# 使用
with DatabaseConnection("localhost", "mydatabase") as conn:
    print("开始数据库操作...")
    print(f"当前连接: {conn}")
    print("数据库操作完成")

contextlib 模块：懒人福音

如果你不想每次都手动写 __enter__ 和 __exit__ 方法，Python还提供了一个 contextlib 模块，里面有一些方便的工具，可以让你更轻松地创建Context Manager。

• contextmanager 装饰器: 这个装饰器可以将一个生成器函数变成一个Context Manager。生成器函数中，yield 之前的代码相当于 __enter__ 方法，yield 之后的代码相当于 __exit__ 方法。

from contextlib import contextmanager

@contextmanager
def file_manager(filename, mode):
    f = None
    try:
        print("准备打开文件...")
        f = open(filename, mode)
        yield f  # 返回文件对象
        print("文件操作完成")
    finally:
        if f:
            print("准备关闭文件...")
            f.close()
            print("文件已关闭")

# 使用
with file_manager("test.txt", "w") as f:
    print("开始写入数据...")
    f.write("Hello, contextlib!")
    print("数据写入完成")

• closing 函数: 这个函数可以将一个具有 close() 方法的对象变成一个Context Manager，保证 close() 方法在退出 with 语句块的时候被调用。

from contextlib import closing
import urllib.request

with closing(urllib.request.urlopen('http://www.example.com')) as webpage:
    for line in webpage:
        print(line)

Context Manager的应用场景

除了文件操作和数据库连接，Context Manager还可以应用于很多其他场景：

• 线程锁: 可以使用Context Manager来管理线程锁，保证锁在使用完毕后一定会被释放，避免死锁。

import threading
from contextlib import contextmanager

lock = threading.Lock()

@contextmanager
def locked():
    print("尝试获取锁...")
    lock.acquire()
    try:
        print("成功获取锁")
        yield
    finally:
        print("释放锁...")
        lock.release()
        print("锁已释放")

# 使用
with locked():
    print("执行需要锁保护的代码...")

• 临时修改环境变量: 可以使用Context Manager来临时修改环境变量，保证在退出 with 语句块的时候环境变量恢复原状。

import os
from contextlib import contextmanager

@contextmanager
def set_environment(key, value):
    original_value = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if original_value is None:
            del os.environ[key]
        else:
            os.environ[key] = original_value

# 使用
with set_environment("MY_VARIABLE", "new_value"):
    print(f"MY_VARIABLE in context: {os.environ['MY_VARIABLE']}")

print(f"MY_VARIABLE after context: {os.environ.get('MY_VARIABLE')}")

表格总结：__enter__ 和 __exit__ 方法

方法名	作用	调用时机	返回值