PYTHON进阶：玩转 with 上下文管理 context manager

2024-11-13

Python 的上下文管理器（Context Manager）是一种用于管理资源的工具，通常用于确保在使用资源时能够正确地进行初始化和清理，例如打开文件、网络连接或数据库连接等。上下文管理器的主要优势是可以自动处理资源的分配和释放，避免资源泄漏。

使用 `with` 语句

上下文管理器通常与 with 语句一起使用。with 语句确保在代码块执行完毕后，资源会被正确释放，无论代码块是正常结束还是由于异常而提前退出。

基本语法

with context_manager [as target]:
    # 在上下文中执行的代码

原理解析

上下文管理器的工作原理基于两个特殊方法：

1. __enter__(): 进入上下文时调用
2. __exit__(exc_type, exc_value, traceback): 退出上下文时调用

当执行 with 语句时，Python 解释器会：

1. 调用上下文管理器的 __enter__() 方法
2. 将 __enter__() 方法的返回值赋值给 as 子句中的变量（如果有）
3. 执行 with 语句体中的代码
4. 调用上下文管理器的 __exit__() 方法

class ExampleContextManager:
    def __init__(self):
        print("初始化上下文管理器")

    def __enter__(self):
        print("进入上下文")
        return self  # 返回值会被赋给 as 子句的变量

    def __exit__(self, exc_type, exc_value, traceback):
        print("退出上下文")
        print(f"异常信息: {exc_type}, {exc_value}, {traceback}")
        return False  # 返回 False 则继续传播异常，True 则抑制异常

# 使用示例
try:
    with ExampleContextManager() as cm:
        print("执行代码")
        raise ValueError("测试异常")
except ValueError:
    print("捕获到异常")

常见应用场景

1. 文件操作

with open('example.txt', 'r') as file:
    content = file.read()
    print(content)
# 文件在这里自动关闭，不需要手动调用 file.close()

2. 数据库连接

with database.connect() as connection:
    cursor = connection.cursor()
    cursor.execute("SELECT * FROM users")
    results = cursor.fetchall()
# 连接会自动关闭

3. 线程锁

import threading

lock = threading.Lock()
with lock:
    # 这里的代码是线程安全的
    # 退出时自动释放锁

4. 临时目录操作

from tempfile import TemporaryDirectory

with TemporaryDirectory() as temp_dir:
    # 在临时目录中进行操作
    print(f"临时目录路径: {temp_dir}")
    # with 语句结束后，临时目录会被自动删除

5. 计时器

from contextlib import contextmanager
import time

@contextmanager
def timer(description):
    start = time.time()
    yield
    elapsed_time = time.time() - start
    print(f"{description}: {elapsed_time:.2f} 秒")

# 使用示例
with timer("处理数据"):
    # 执行一些耗时操作
    time.sleep(1)

6. 数据库事务管理

from contextlib import contextmanager
import sqlite3

@contextmanager
def transaction(connection):
    cursor = connection.cursor()
    try:
        yield cursor
        connection.commit()
    except Exception:
        connection.rollback()
        raise

# 使用示例
with sqlite3.connect('database.db') as conn:
    with transaction(conn) as cursor:
        cursor.execute("INSERT INTO users (name) VALUES (?)", ("张三",))
        cursor.execute("UPDATE users SET age = ? WHERE name = ?", (25, "张三"))

7. 网络连接管理

import socket
from contextlib import contextmanager

@contextmanager
def tcp_connection(host, port):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        sock.connect((host, port))
        yield sock
    finally:
        sock.close()

# 使用示例
with tcp_connection("example.com", 80) as sock:
    sock.send(b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
    response = sock.recv(1024)

8. 环境变量临时修改

import os
from contextlib import contextmanager

@contextmanager
def environ_temp(**kwargs):
    old_values = {}
    try:
        # 保存旧值
        for key, value in kwargs.items():
            old_values[key] = os.environ.get(key)
            os.environ[key] = value
        yield
    finally:
        # 恢复旧值
        for key in kwargs:
            if old_values[key] is None:
                del os.environ[key]
            else:
                os.environ[key] = old_values[key]

# 使用示例
with environ_temp(DEBUG='1', ENV='testing'):
    print(os.environ['DEBUG'])  # 输出: 1
# 环境变量在这里恢复原值

自定义上下文管理器

使用类实现

class MyContextManager:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"进入上下文: {self.name}")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"退出上下文: {self.name}")
        if exc_type:
            print(f"异常类型: {exc_type}, 异常值: {exc_value}")
        # 返回 True 可以抑制异常
        return False

# 使用示例
with MyContextManager("测试") as manager:
    print("在上下文中执行代码")
    # raise ValueError("测试异常")

使用装饰器实现

from contextlib import contextmanager

@contextmanager
def my_context_manager(name):
    print(f"进入上下文: {name}")
    try:
        yield
    finally:
        print(f"退出上下文: {name}")

# 使用示例
with my_context_manager("测试"):
    print("在上下文中执行代码")

contextlib 模块

Python 的 contextlib 模块提供了一些有用的工具：

1. @contextmanager 装饰器：简化上下文管理器的创建
2. closing()：自动调用对象的 close() 方法
3. suppress()：忽略特定异常
4. nullcontext()：空的上下文管理器
5. ExitStack：动态管理一组上下文管理器

from contextlib import ExitStack

def process_files(file_list):
    with ExitStack() as stack:
        files = [stack.enter_context(open(fname)) for fname in file_list]
        # 现在所有文件都打开了，可以进行处理
        for file in files:
            print(file.read())
        # 退出 with 语句后，所有文件都会自动关闭

高级用法

1. 嵌套使用

with context_manager1() as ctx1:
    with context_manager2() as ctx2:
        # 使用 ctx1 和 ctx2
        pass

# 或者使用更简洁的写法
with context_manager1() as ctx1, context_manager2() as ctx2:
    # 使用 ctx1 和 ctx2
    pass

2. 异步上下文管理器

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

    async def __aexit__(self, exc_type, exc_value, traceback):
        print("退出异步上下文")
        return False

# 使用异步上下文管理器
async def main():
    async with AsyncContextManager():
        print("在异步上下文中执行代码")

3. 可重入的上下文管理器

class ReentrantContextManager:
    def __init__(self):
        self._lock = threading.RLock()

    def __enter__(self):
        self._lock.acquire()
        return self

    def __exit__(self, *args):
        self._lock.release()

# 可以多次进入而不会死锁
with ReentrantContextManager() as rcm:
    with rcm:
        print("嵌套使用相同的上下文管理器")

常见陷阱和注意事项

1. 资源泄露问题

def wrong_usage():
    # 错误：没有使用 with 语句
    file = open('example.txt')
    # 如果这里发生异常，文件句柄可能永远不会被关闭

def correct_usage():
    # 正确：使用 with 语句
    with open('example.txt') as file:
        # 文件会被自动关闭，即使发生异常
        pass

2. 异常处理

class DatabaseConnection:
    def __enter__(self):
        # 建立数据库连接
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # 确保在发生异常时也能正确清理资源
        try:
            # 清理资源
            self.close()
        except Exception as e:
            # 记录清理过程中的错误，但不抑制原始异常
            print(f"清理资源时发生错误: {e}")
            return False