Python 自定义上下文管理器

创建完整的上下文管理器类实现资源管理，支持可重入、异常安全等高级特性。

基础实现

class FileManager:
    "文件管理上下文管理器"

    def __init__(self, filename, mode='r'):
        self.filename = filename
        self.mode = mode
        self.file = None

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            self.file.close()
        return False

# 使用
with FileManager('data.txt', 'w') as f:
    f.write('Hello, World!')

计时器上下文管理器

import time

class Timer:
    "计时器上下文管理器"

    def __init__(self, name=None):
        self.name = name
        self.start = None
        self.elapsed = None

    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        prefix = f"{self.name}: " if self.name else ""
        print(f"{prefix}耗时 {self.elapsed:.4f} 秒")
        return False

# 使用
with Timer("数据处理"):
    sum(range(1000000))
# 输出: 数据处理: 耗时 0.0xxx 秒

数据库连接管理器

import sqlite3

class DatabaseConnection:
    "数据库连接上下文管理器"

    def __init__(self, db_path):
        self.db_path = db_path
        self.conn = None
        self.cursor = None

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_path)
        self.cursor = self.conn.cursor()
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.conn.commit()  # 无异常则提交
        else:
            self.conn.rollback()  # 有异常则回滚

        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()
        return False

# 使用
with DatabaseConnection('test.db') as cursor:
    cursor.execute('CREATE TABLE IF NOT EXISTS users (id INTEGER, name TEXT)')
    cursor.execute('INSERT INTO users VALUES (1, "Alice")')

可重入上下文管理器

class ReentrantCounter:
    "可重入的计数器上下文管理器"

    def __init__(self):
        self._count = 0
        self._depth = 0

    def __enter__(self):
        self._depth += 1
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._depth -= 1
        if self._depth == 0 and exc_type is None:
            self._count += 1
        return False

    @property
    def count(self):
        return self._count

# 使用
counter = ReentrantCounter()
with counter:
    with counter:  # 可嵌套使用
        print(f"深度: {counter._depth}")  # 2

print(f"计数: {counter.count}")  # 1

状态保存与恢复

class WorkingDirectory:
    "临时切换工作目录"

    def __init__(self, new_dir):
        self.new_dir = new_dir
        self.old_dir = None

    def __enter__(self):
        import os
        self.old_dir = os.getcwd()
        os.chdir(self.new_dir)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        import os
        os.chdir(self.old_dir)
        return False

# 使用
with WorkingDirectory('/tmp'):
    # 当前工作目录为 /tmp
    pass
# 自动恢复原工作目录

抽象基类继承

from abc import ABC, abstractmethod

class BaseContextManager(ABC):
    "上下文管理器抽象基类"

    def __enter__(self):
        self.setup()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return self.cleanup(exc_type, exc_val, exc_tb)

    @abstractmethod
    def setup(self):
        "子类实现资源获取"
        pass

    @abstractmethod
    def cleanup(self, exc_type, exc_val, exc_tb):
        "子类实现资源清理"
        pass

class MyContext(BaseContextManager):
    def setup(self):
        print("资源获取")

    def cleanup(self, exc_type, exc_val, exc_tb):
        print("资源清理")
        return False

设计原则

要点总结

• __enter__ 获取资源并返回，__exit__ 释放资源
• 异常处理应在 __exit__ 中完成
• 支持事务模式：成功提交，失败回滚
• 可重入设计需跟踪嵌套深度
• 抽象基类可规范上下文管理器接口
• 适用于文件、连接、锁、计时等资源管理