Python 上下文管理器使用技巧四个实例

什么是上下文管理器？

上下文管理器是一种用于设置和清理资源的协议。它通常用于 with 语句中，确保在进入和退出代码块时执行特定的操作。上下文管理器的主要优点是它可以自动管理资源，避免资源泄漏。

基本语法

上下文管理器的基本语法如下：

with context_manager_expression as variable:
    # 代码块

其中，context_manager_expression 是一个实现了上下文管理器协议的对象，variable 是可选的，用于接收上下文管理器返回的值。

实现上下文管理器的方法

上下文管理器可以通过两种方式实现：

• 类方法：通过实现 __enter__ 和 __exit__ 方法。
• 装饰器方法：使用 contextlib 模块中的 contextmanager 装饰器。

示例 1：文件操作

文件操作是最常见的使用上下文管理器的场景之一。我们可以使用 with 语句来打开文件，确保文件在使用后自动关闭。

# 使用上下文管理器打开文件
with open('example.txt', 'w') as file:
    file.write('Hello, World!')

# 文件在离开 with 代码块后自动关闭
print(file.closed)  # 输出: True

示例 2：数据库连接

在处理数据库连接时，上下文管理器可以确保在操作完成后自动关闭连接，避免资源泄漏。

import sqlite3

# 使用上下文管理器管理数据库连接
with sqlite3.connect('example.db') as conn:
    cursor = conn.cursor()
    cursor.execute('CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)')
    cursor.execute('INSERT INTO users (name) VALUES (?)', ('Alice',))
    conn.commit()

# 连接在离开 with 代码块后自动关闭
print(conn.in_transaction)  # 输出: False

示例 3：自定义上下文管理器（类方法）

我们可以自定义一个上下文管理器来管理特定的资源。例如，假设我们需要一个上下文管理器来记录代码块的执行时间。

import time

class Timer:
    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.time()
        print(f'Time taken: {self.end_time - self.start_time:.2f} seconds')

# 使用自定义的上下文管理器
with Timer() as timer:
    time.sleep(2)

# 输出: Time taken: 2.00 seconds

示例 4：自定义上下文管理器（装饰器方法）

使用 contextlib 模块中的 contextmanager 装饰器可以更简洁地实现上下文管理器。

from contextlib import contextmanager

@contextmanager
def managed_resource():
    print('Resource acquired')
    try:
        yield 'resource'
    finally:
        print('Resource released')

# 使用装饰器实现的上下文管理器
with managed_resource() as resource:
    print(f'Using {resource}')

# 输出:
# Resource acquired
# Using resource
# Resource released

实战案例：管理多个文件的读写操作

假设我们需要同时读取多个文件并将其内容合并到一个新文件中。我们可以使用上下文管理器来确保所有文件在操作完成后都能正确关闭。

def merge_files(input_files, output_file):
    with open(output_file, 'w') as outfile:
        for input_file in input_files:
            with open(input_file, 'r') as infile:
                outfile.write(infile.read())
                outfile.write('\n')

# 测试数据
input_files = ['file1.txt', 'file2.txt']
output_file = 'merged.txt'

# 写入测试数据
with open('file1.txt', 'w') as f1:
    f1.write('Content of file 1')
with open('file2.txt', 'w') as f2:
    f2.write('Content of file 2')

# 合并文件
merge_files(input_files, output_file)

# 验证结果
with open('merged.txt', 'r') as merged:
    print(merged.read())

# 输出:
# Content of file 1
# Content of file 2

总结

本文介绍了 Python 上下文管理器的基本概念和实现方法，并通过四个实例详细展示了如何在不同场景中使用上下文管理器。