Python 上下文管理器协议：`__enter__`, `__exit__` 的高级用法

各位观众，掌声在哪里！今天咱们来聊聊Python里一个听起来高大上，用起来贼顺手的玩意儿：上下文管理器。别怕，这名字唬人，其实就是个负责任的好管家，帮你自动搞定一些收尾工作。咱们今天不光要了解它，还要深入挖掘它的高级用法，保证让各位看完之后，觉得这玩意儿真香！

什么是上下文管理器？（别告诉我你只知道with）

首先，别听到“上下文管理器”就觉得头大。简单来说，它就是一个对象，定义了在使用with语句时，进入和退出代码块时需要执行的操作。最常见的例子就是文件操作：

with open("my_file.txt", "w") as f:
    f.write("Hello, world!")
# 文件会自动关闭，不用你操心

这里，open()函数返回的对象就是一个上下文管理器。with语句负责在进入代码块之前调用__enter__方法，在退出代码块之后调用__exit__方法。这样，文件打开和关闭的操作就被自动管理起来了，再也不用担心忘记关闭文件导致资源泄露了！

__enter__和__exit__：幕后英雄

要理解上下文管理器的核心，就得搞清楚__enter__和__exit__这两个方法。

• __enter__(self): 这个方法在进入with语句块之前被调用。它通常负责准备资源，比如打开文件、建立数据库连接、获取锁等等。它必须返回一个值，这个值会被赋值给with语句中的as后面的变量（如果没有as，则忽略返回值）。

• __exit__(self, exc_type, exc_val, exc_tb): 这个方法在退出with语句块之后被调用。无论代码块是正常执行完毕，还是发生了异常，它都会被执行。它负责清理资源，比如关闭文件、释放锁、回滚事务等等。

  • exc_type: 异常类型 (如果代码块中没有发生异常，则为None)
  • exc_val: 异常实例 (如果代码块中没有发生异常，则为None)
  • exc_tb: 异常回溯对象 (如果代码块中没有发生异常，则为None)

  __exit__方法应该返回一个布尔值：

  • True: 表示该方法已经处理了异常，异常不会被传播到with语句块之外。
  • False (或None): 表示该方法没有处理异常，异常会被传播到with语句块之外。

自己动手，丰衣足食：自定义上下文管理器

光会用别人写好的上下文管理器可不行，咱们得学会自己写！这样才能真正掌握它的精髓。

class MyContextManager:
    def __init__(self, resource_name):
        self.resource_name = resource_name
        self.resource = None

    def __enter__(self):
        print(f"准备获取资源: {self.resource_name}")
        self.resource = open(self.resource_name, "w") # 假设这里是打开文件
        return self.resource

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"清理资源: {self.resource_name}")
        if self.resource:
            self.resource.close()
        if exc_type:
            print(f"检测到异常: {exc_type}, {exc_val}")
            # 可以选择处理异常，比如记录日志
            return True  # 返回True表示已处理异常，阻止异常传播
        return False # 返回False表示没有处理异常，异常继续传播

# 使用自定义上下文管理器
with MyContextManager("my_custom_file.txt") as f:
    f.write("This is from my custom context manager!n")
    # raise ValueError("Something went wrong!") # 取消注释可以模拟异常情况

print("代码执行完毕！")

在这个例子中，我们定义了一个MyContextManager类，它模拟了打开和关闭文件的操作。__enter__方法负责打开文件，并返回文件对象。__exit__方法负责关闭文件，并处理可能发生的异常。

高级用法一：处理异常

__exit__方法的一个重要作用就是处理异常。我们可以根据exc_type、exc_val和exc_tb来判断是否发生了异常，以及异常的类型和信息。

class DatabaseTransaction:
    def __init__(self, db_connection):
        self.db_connection = db_connection
        self.cursor = None

    def __enter__(self):
        self.cursor = self.db_connection.cursor()
        self.cursor.execute("START TRANSACTION")
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            print("事务回滚！")
            self.db_connection.rollback()
        else:
            print("事务提交！")
            self.db_connection.commit()
        self.cursor.close()
        return False # 别忘了返回False，让异常继续传播（如果存在）

# 使用数据库事务
import sqlite3
conn = sqlite3.connect(':memory:')  # 使用内存数据库进行测试

try:
    with DatabaseTransaction(conn) as cursor:
        cursor.execute("CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)")
        cursor.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
        cursor.execute("INSERT INTO users (name) VALUES (?)", ("Bob",))
        # 模拟一个异常
        cursor.execute("SELECT * FROM non_existent_table") # 会抛出OperationalError
except sqlite3.Error as e:
    print(f"捕获到异常: {e}")

conn.close()

在这个例子中，我们模拟了一个数据库事务。__enter__方法负责启动事务，__exit__方法负责提交或回滚事务，并关闭游标。如果代码块中发生了异常，__exit__方法会回滚事务，保证数据的一致性。

高级用法二：重入的上下文管理器

有时候，你可能需要在同一个with语句块中多次进入同一个上下文管理器。这被称为重入。要实现重入，你需要保证__enter__方法在每次调用时都能正确地准备资源，并且__exit__方法在每次调用时都能正确地清理资源。

class ReentrantContextManager:
    def __init__(self, name):
        self.name = name
        self.enter_count = 0

    def __enter__(self):
        self.enter_count += 1
        print(f"进入上下文 {self.name}, 第 {self.enter_count} 次")
        return self  # 返回self，方便在with语句中使用

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"退出上下文 {self.name}, 第 {self.enter_count} 次")
        self.enter_count -= 1
        return False

# 使用重入的上下文管理器
with ReentrantContextManager("MyReentrantContext") as context:
    print("第一次进入上下文")
    with context:
        print("第二次进入上下文")
    print("第一次退出上下文")
print("完全退出上下文")

在这个例子中，__enter__方法简单地增加了一个计数器，并打印一条消息。__exit__方法减少计数器，并打印一条消息。这样，我们就可以清楚地看到上下文管理器被进入和退出的次数。

高级用法三：使用contextlib简化代码

Python的contextlib模块提供了一些工具，可以简化上下文管理器的编写。最常用的工具是contextmanager装饰器。

from contextlib import contextmanager

@contextmanager
def my_context_manager(resource_name):
    print(f"准备获取资源: {resource_name}")
    resource = open(resource_name, "w")
    try:
        yield resource # 关键：yield语句返回资源
    finally:
        print(f"清理资源: {resource_name}")
        resource.close()

# 使用@contextmanager装饰器创建的上下文管理器
with my_context_manager("my_contextlib_file.txt") as f:
    f.write("This is from contextlib!n")

使用contextmanager装饰器，我们可以用一个生成器函数来定义上下文管理器。yield语句将函数分成两部分：yield之前的代码相当于__enter__方法，yield之后的代码相当于__exit__方法。yield返回的值会被赋值给with语句中的as后面的变量。

表格总结：__enter__ vs. __exit__

特性	__enter__	__exit__
调用时机	进入with语句块之前	退出with语句块之后
主要作用	准备资源	清理资源，处理异常
参数	self	self, exc_type, exc_val, exc_tb
返回值	返回一个值，赋值给with语句的as变量	布尔值：True表示已处理异常，False或None表示未处理异常
是否必须实现	必须实现	必须实现

一些额外的思考 (敲黑板!)

• 上下文管理器并非只能用于资源管理。 它可以用于任何需要在进入和退出代码块时执行特定操作的场景，例如性能分析、状态管理、安全检查等等。

• 不要在__enter__方法中做过多的事情。 __enter__方法应该尽可能地简单和快速，避免阻塞主线程。

• __exit__方法应该总是被执行。 即使__enter__方法抛出了异常，__exit__方法也应该被执行，以保证资源的正确清理。

• 异常处理的策略要根据具体情况来决定。 有时候，我们希望捕获并处理异常，阻止异常传播；有时候，我们希望让异常继续传播，以便上层代码能够处理。

案例分析：一个更复杂的例子 – 线程锁

咱们来一个更实际的例子，使用上下文管理器管理线程锁。

import threading
import time

class ThreadSafeCounter:
    def __init__(self):
        self.count = 0
        self.lock = threading.Lock()

    def increment(self):
        with self.lock: # 使用锁作为上下文管理器
            self.count += 1

    def get_count(self):
        with self.lock:
            return self.count

# 测试代码
counter = ThreadSafeCounter()

def worker():
    for _ in range(100000):
        counter.increment()

threads = []
for _ in range(5):
    thread = threading.Thread(target=worker)
    threads.append(thread)
    thread.start()

for thread in threads:
    thread.join()

print(f"最终计数: {counter.get_count()}")

在这个例子中，threading.Lock本身就是一个上下文管理器。当我们使用with self.lock:时，__enter__方法会获取锁，__exit__方法会释放锁。这样，我们就可以保证在访问count变量时，只有一个线程可以访问，避免了竞态条件。

更进一步：异步上下文管理器 (Async Context Managers)

如果你在使用asyncio，那么你肯定会遇到异步上下文管理器。它们和普通的上下文管理器类似，但是使用async关键字定义。

import asyncio

class AsyncContextManager:
    async def __aenter__(self):
        print("进入异步上下文")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        print("退出异步上下文")

async def main():
    async with AsyncContextManager():
        print("在异步上下文中")

asyncio.run(main())

这里，__aenter__和__aexit__方法都是协程函数，需要使用await关键字来调用。

总结：上下文管理器，你的代码管家！

好了，各位，今天的讲座就到这里。希望通过今天的学习，大家对Python的上下文管理器有了更深入的理解。记住，上下文管理器不仅仅是with语句的语法糖，更是一种优雅的代码组织方式，可以帮助我们编写更健壮、更易于维护的代码。 以后写代码，记得带上你的代码管家！

感谢各位的观看，下次再见！