第80集：上下文管理器

学习目标

• 理解上下文管理器的概念
• 掌握with语句的使用
• 学会创建自定义上下文管理器
• 理解__enter__和__exit__方法
• 掌握contextlib模块的使用
• 了解上下文管理器的应用场景

什么是上下文管理器

上下文管理器是一种Python对象，它定义了在进入和退出上下文时应该执行的代码。上下文管理器通常与with语句一起使用，用于自动管理资源，确保资源在使用后被正确释放。

上下文管理器的特点

• 自动管理资源
• 确保资源被正确释放
• 代码更简洁安全
• 支持异常处理

with语句基础

基本语法

with context_manager as variable:
    # 使用资源的代码
    pass

示例1：文件操作

# 使用with语句打开文件
with open('file.txt', 'r', encoding='utf-8') as file:
    content = file.read()
    print(content)
# 文件会自动关闭，无需手动调用close()

示例2：锁的使用

import threading

lock = threading.Lock()

with lock:
    # 临界区代码
    print("执行临界区操作")
# 锁会自动释放

自定义上下文管理器

基本结构

class MyContextManager:
    def __enter__(self):
        # 进入上下文时执行
        pass
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 退出上下文时执行
        pass
        return False

示例3：简单上下文管理器

class SimpleContext:
    """简单上下文管理器"""
    def __enter__(self):
        print("进入上下文")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出上下文")
        return False

with SimpleContext():
    print("在上下文中执行代码")

示例4：计时上下文管理器

import time

class Timer:
    """计时上下文管理器"""
    def __enter__(self):
        self.start = time.time()
        print("开始计时")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        elapsed = time.time() - self.start
        print(f"耗时: {elapsed:.4f}秒")
        return False

with Timer():
    time.sleep(1)
    print("执行一些操作")

示例5：资源管理上下文管理器

class ResourceManager:
    """资源管理上下文管理器"""
    def __init__(self, resource_name):
        self.resource_name = resource_name
        self.resource = None
    
    def __enter__(self):
        print(f"获取资源: {self.resource_name}")
        self.resource = f"资源_{self.resource_name}"
        return self.resource
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"释放资源: {self.resource_name}")
        self.resource = None
        return False

with ResourceManager("database") as resource:
    print(f"使用资源: {resource}")

__enter__和__exit__方法详解

__enter__方法

def __enter__(self):
    """
    进入上下文时调用
    
    Returns:
        返回的对象会赋值给as后面的变量
    """
    pass

__exit__方法

def __exit__(self, exc_type, exc_val, exc_tb):
    """
    退出上下文时调用
    
    Args:
        exc_type: 异常类型，如果没有异常则为None
        exc_val: 异常值，如果没有异常则为None
        exc_tb: 异常追踪信息，如果没有异常则为None
    
    Returns:
        True: 抑制异常
        False: 不抑制异常，继续传播
    """
    pass

示例6：异常处理

class ExceptionHandler:
    """异常处理上下文管理器"""
    def __enter__(self):
        print("进入上下文")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print(f"捕获到异常: {exc_type.__name__}")
            print(f"异常信息: {exc_val}")
            return True  # 抑制异常
        print("正常退出上下文")
        return False

with ExceptionHandler():
    print("执行代码")
    raise ValueError("这是一个测试异常")

print("程序继续执行")

示例7：不抑制异常

class NoSuppressException:
    """不抑制异常的上下文管理器"""
    def __enter__(self):
        print("进入上下文")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print(f"捕获到异常: {exc_type.__name__}")
            print(f"异常信息: {exc_val}")
            return False  # 不抑制异常
        print("正常退出上下文")
        return False

try:
    with NoSuppressException():
        print("执行代码")
        raise ValueError("这是一个测试异常")
except ValueError as e:
    print(f"异常被重新抛出: {e}")

contextlib模块

contextlib.contextmanager装饰器

from contextlib import contextmanager

@contextmanager
def simple_context():
    """使用装饰器创建上下文管理器"""
    print("进入上下文")
    try:
        yield  # yield之前是__enter__的代码
    finally:
        print("退出上下文")  # yield之后是__exit__的代码

with simple_context():
    print("在上下文中执行")

示例8：使用contextmanager创建计时器

from contextlib import contextmanager
import time

@contextmanager
def timer():
    """计时上下文管理器"""
    start = time.time()
    print("开始计时")
    try:
        yield start
    finally:
        elapsed = time.time() - start
        print(f"耗时: {elapsed:.4f}秒")

with timer():
    time.sleep(0.5)
    print("执行操作")

示例9：使用contextmanager创建资源管理器

from contextlib import contextmanager

@contextmanager
def resource_manager(resource_name):
    """资源管理上下文管理器"""
    print(f"获取资源: {resource_name}")
    resource = f"资源_{resource_name}"
    try:
        yield resource
    finally:
        print(f"释放资源: {resource_name}")

with resource_manager("database") as resource:
    print(f"使用资源: {resource}")

示例10：使用contextmanager处理异常

from contextlib import contextmanager

@contextmanager
def exception_handler():
    """异常处理上下文管理器"""
    print("进入上下文")
    try:
        yield
    except Exception as e:
        print(f"捕获到异常: {e}")
        raise  # 重新抛出异常
    finally:
        print("退出上下文")

try:
    with exception_handler():
        print("执行代码")
        raise ValueError("测试异常")
except ValueError:
    print("异常被重新抛出")

contextlib的其他工具

contextlib.closing

from contextlib import closing

class Resource:
    """需要关闭的资源"""
    def close(self):
        print("资源已关闭")

with closing(Resource()):
    print("使用资源")
# 自动调用close()方法

contextlib.suppress

from contextlib import suppress

# 抑制特定异常
with suppress(FileNotFoundError):
    with open('不存在的文件.txt', 'r') as f:
        content = f.read()

print("程序继续执行")

contextlib.redirect_stdout

from contextlib import redirect_stdout
import io

# 重定向标准输出
f = io.StringIO()
with redirect_stdout(f):
    print("这行输出会被重定向")

print(f"捕获的输出: {f.getvalue()}")

contextlib.redirect_stderr

from contextlib import redirect_stderr
import io

# 重定向标准错误
f = io.StringIO()
with redirect_stderr(f):
    print("这行错误输出会被重定向", file=sys.stderr)

print(f"捕获的错误输出: {f.getvalue()}")

上下文管理器的应用场景

应用1：数据库连接管理

class DatabaseConnection:
    """数据库连接上下文管理器"""
    def __init__(self, database_url):
        self.database_url = database_url
        self.connection = None
    
    def __enter__(self):
        print(f"连接到数据库: {self.database_url}")
        self.connection = f"连接对象_{self.database_url}"
        return self.connection
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print("发生异常，回滚事务")
        else:
            print("提交事务")
        print("关闭数据库连接")
        self.connection = None
        return False

with DatabaseConnection("mysql://localhost/mydb") as conn:
    print(f"执行查询: {conn}")

应用2：临时目录管理

import tempfile
import os

class TemporaryDirectory:
    """临时目录上下文管理器"""
    def __init__(self):
        self.temp_dir = None
    
    def __enter__(self):
        self.temp_dir = tempfile.mkdtemp()
        print(f"创建临时目录: {self.temp_dir}")
        return self.temp_dir
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.temp_dir and os.path.exists(self.temp_dir):
            print(f"删除临时目录: {self.temp_dir}")
            for root, dirs, files in os.walk(self.temp_dir, topdown=False):
                for name in files:
                    os.remove(os.path.join(root, name))
                for name in dirs:
                    os.rmdir(os.path.join(root, name))
            os.rmdir(self.temp_dir)
        return False

with TemporaryDirectory() as temp_dir:
    print(f"使用临时目录: {temp_dir}")

应用3：临时改变工作目录

import os

class ChangeDirectory:
    """改变工作目录上下文管理器"""
    def __init__(self, new_path):
        self.new_path = new_path
        self.old_path = None
    
    def __enter__(self):
        self.old_path = os.getcwd()
        print(f"切换目录: {self.old_path} -> {self.new_path}")
        os.chdir(self.new_path)
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"恢复目录: {self.new_path} -> {self.old_path}")
        os.chdir(self.old_path)
        return False

with ChangeDirectory("C:/"):
    print(f"当前目录: {os.getcwd()}")

应用4：临时环境变量

import os

class TemporaryEnvironment:
    """临时环境变量上下文管理器"""
    def __init__(self, **kwargs):
        self.env_vars = kwargs
        self.old_values = {}
    
    def __enter__(self):
        for key, value in self.env_vars.items():
            self.old_values[key] = os.environ.get(key)
            os.environ[key] = str(value)
            print(f"设置环境变量: {key}={value}")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        for key, value in self.old_values.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
            print(f"恢复环境变量: {key}={value}")
        return False

with TemporaryEnvironment(DEBUG="True", MODE="test"):
    print(f"DEBUG: {os.environ.get('DEBUG')}")
    print(f"MODE: {os.environ.get('MODE')}")

应用5：性能分析

import time
import sys

class PerformanceProfiler:
    """性能分析上下文管理器"""
    def __init__(self, name="操作"):
        self.name = name
        self.start_time = None
    
    def __enter__(self):
        self.start_time = time.time()
        print(f"开始分析: {self.name}")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        elapsed = time.time() - self.start_time
        print(f"完成分析: {self.name}")
        print(f"执行时间: {elapsed:.4f}秒")
        print(f"内存使用: {sys.getsizeof(self)} 字节")
        return False

with PerformanceProfiler("数据处理"):
    data = [i ** 2 for i in range(100000)]
    sum(data)

嵌套上下文管理器

示例11：嵌套with语句

class Context1:
    def __enter__(self):
        print("进入Context1")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出Context1")
        return False

class Context2:
    def __enter__(self):
        print("进入Context2")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退出Context2")
        return False

# 嵌套使用
with Context1():
    with Context2():
        print("在嵌套上下文中执行")

示例12：使用多个上下文管理器

# 使用多个上下文管理器
with Context1(), Context2():
    print("使用多个上下文管理器")

上下文管理器的限制

限制1：不能在异步函数中使用

# 不能在异步函数中使用普通上下文管理器
async def async_function():
    with Context1():  # 错误
        pass

# 需要使用异步上下文管理器
class AsyncContext:
    async def __aenter__(self):
        pass
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass

限制2：__exit__方法的返回值

# __exit__方法的返回值很重要
class BadContext:
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 忘记返回False，默认返回None，相当于False
        pass

class GoodContext:
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 明确返回False
        return False

最佳实践

实践1：确保资源释放

class SafeResource:
    """安全的资源管理器"""
    def __init__(self, resource_name):
        self.resource_name = resource_name
        self.resource = None
    
    def __enter__(self):
        self.resource = self.acquire_resource()
        return self.resource
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            if self.resource is not None:
                self.release_resource()
        except Exception as e:
            print(f"释放资源时出错: {e}")
        return False
    
    def acquire_resource(self):
        print(f"获取资源: {self.resource_name}")
        return f"资源_{self.resource_name}"
    
    def release_resource(self):
        print(f"释放资源: {self.resource_name}")
        self.resource = None

实践2：使用contextmanager简化代码

from contextlib import contextmanager

@contextmanager
def managed_resource(resource_name):
    """使用contextmanager简化代码"""
    resource = None
    try:
        resource = acquire_resource(resource_name)
        yield resource
    finally:
        if resource is not None:
            release_resource(resource_name)

def acquire_resource(name):
    print(f"获取资源: {name}")
    return f"资源_{name}"

def release_resource(name):
    print(f"释放资源: {name}")

with managed_resource("database") as resource:
    print(f"使用资源: {resource}")

实践3：清晰的命名

class DatabaseConnectionContext:
    """清晰的命名"""
    def __init__(self, connection_string):
        self.connection_string = connection_string
    
    def __enter__(self):
        pass
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

实践4：文档化上下文管理器

class TimerContext:
    """
    计时上下文管理器
    
    用于测量代码块的执行时间
    
    Attributes:
        name: 操作名称
    """
    def __init__(self, name="操作"):
        self.name = name
        self.start_time = None
    
    def __enter__(self):
        import time
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        import time
        elapsed = time.time() - self.start_time
        print(f"{self.name} 耗时: {elapsed:.4f}秒")
        return False

常见错误

错误1：忘记实现__exit__方法

# 错误：忘记实现__exit__方法
class BadContext:
    def __enter__(self):
        return self

# 正确：实现__exit__方法
class GoodContext:
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        return False

错误2：__exit__方法返回True抑制异常

# 错误：意外抑制异常
class BadContext:
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        return True  # 抑制所有异常

# 正确：根据需要决定是否抑制异常
class GoodContext:
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print(f"处理异常: {exc_val}")
            return False  # 不抑制异常
        return False

错误3：资源未正确释放

# 错误：资源未正确释放
class BadContext:
    def __init__(self):
        self.resource = None
    
    def __enter__(self):
        self.resource = acquire_resource()
        return self.resource
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 忘记释放资源
        pass

# 正确：确保资源释放
class GoodContext:
    def __init__(self):
        self.resource = None
    
    def __enter__(self):
        self.resource = acquire_resource()
        return self.resource
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            if self.resource is not None:
                release_resource()
        except Exception as e:
            print(f"释放资源时出错: {e}")
        return False

总结

上下文管理器是Python中强大的工具，具有以下特点：

核心概念

• 上下文管理器通过实现__enter__和__exit__方法实现
• 与with语句一起使用，自动管理资源
• 可以使用contextlib.contextmanager装饰器简化实现

优势

• 自动管理资源
• 代码更简洁安全
• 支持异常处理
• 避免资源泄漏

应用场景

• 文件操作
• 数据库连接
• 锁管理
• 临时目录
• 性能分析
• 环境变量管理

掌握上下文管理器将帮助你编写更安全、更优雅的Python代码。上下文管理器是Python中重要的概念，理解它的工作原理对于编写高质量的Python程序至关重要。