在 Python 中编写自定义上下文管理器

动机

您也许已经熟悉 with 语句——这是一种在 Python 中简洁处理资源的方法。但您是否想过它在底层是如何工作的？实际上，with 关键字的威力来自上下文管理器。

上下文管理器是 Python 中一种基础的设计模式，它为资源管理提供了结构化的方法。它确保资源被获取、被正确使用，并最终被释放或清理，即使发生错误或异常也不例外。

当处理文件、网络连接或数据库句柄等资源时，这一点尤为重要。借助上下文管理器，我们可以编写更简洁、更可靠的代码，不必再担心忘记关闭文件或释放锁。

在本教程中，我们将超越 Python 内置的上下文管理器，学习如何编写自己的管理器。

理解 Python 中的上下文管理器

在底层，上下文管理器是定义了两个特殊方法的对象：__enter__ 和 __exit__。当您进入 with 代码块时会调用 __enter__ 方法，其返回值会被赋给该块内的变量。另一方面，无论 with 代码块是正常结束还是因异常退出，__exit__ 方法都会被调用。

这种结构确保了资源被正确处理。我们来看一些 Python 内置的上下文管理器来加深理解：

1. 文件管理

以打开文件这个经典示例为例：

with open('file.txt', 'r') as file:
   data = file.read()
技术支持

这里，open('file.txt', 'r') 充当上下文管理器。当您进入 with 代码块时，会调用文件对象的 enter 方法，打开文件并将其赋给变量 file。随后，您可以使用 file.read() 访问文件内容。

接着，当 with 代码块退出时，即使发生异常，也会保证调用文件对象的 __exit__ 方法。该方法负责关闭文件，确保不会遗留未关闭的文件句柄。

2. 线程加锁

除了文件之外，上下文管理器还可用于通过 threading.Lock() 进行线程同步：

import threading

lock = threading.Lock()
with lock:
   # Critical section
   print("This code is executed under lock protection.")
技术支持

这里，lock 是一个 threading.Lock 对象，也是一个上下文管理器。进入 with 代码块时，__enter__ 方法会获取锁，确保同一时间只有一个线程能执行临界区。退出 with 代码块时，__exit__ 方法会释放锁，使其他线程得以继续。

3. 数据库连接

类似地，上下文管理器可用于管理数据库连接：

import sqlite3

with sqlite3.connect('database.db') as connection:
    cursor = connection.cursor()
    cursor.execute("SELECT * FROM table")
    rows = cursor.fetchall()
技术支持

sqlite3.connect('database.db') 调用返回一个上下文管理器。进入 with 代码块时会建立到数据库的连接，并将其赋给 connection。随后您可以使用 cursor 与数据库交互。__exit__ 方法保证在 with 代码块退出时关闭连接，防止资源泄漏。

4. 网络套接字

上下文管理器甚至可以处理网络通信：

import socket

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.connect(('localhost', 8080))
    s.sendall(b'Hello, world')
    data = s.recv(1024)
技术支持

这里，socket.socket(socket.AF_INET, socket.SOCK_STREAM) 创建了一个作为上下文管理器的套接字对象。__enter__ 方法创建套接字，在 with 代码块内，您可以连接、发送数据并接收数据。__exit__ 方法确保在结束时正确关闭套接字。

5. 目录扫描

这一次，os.scandir('.') 提供了遍历目录条目的方式：

with os.scandir('.') as entries:
   for entry in entries:
       print(entry.name)
技术支持

os.scandir('.') 在这里充当上下文管理器。__enter__ 方法开启目录扫描，您可以在 with 代码块中遍历条目。__exit__ 方法在退出时清理目录扫描。

可见，由上下文管理器驱动的 with 语句通过自动处理资源的分配与释放，大大简化了资源管理。

在 Python 中编写自定义上下文管理器

现在，让我们看看如何编写自定义管理器，从而对资源管理进行更精细的控制。

编写自定义上下文管理器主要有两种方式：基于类和基于函数。

基于类的方式

基于类的方式是编写上下文管理器最具结构性、也最灵活的方法。在这里，您定义一个类，并实现特殊方法 __enter__ 和 __exit__。

我们来看看一个用于衡量执行时间的 Timer 类示例：

import time


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

   def __exit__(self, exc_type, exc_value, traceback):
       self.end_time = time.time()
       elapsed_time = self.end_time - self.start_time
       print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with Timer() as timer:
       # Code block to measure the execution time
       time.sleep(2)  # Simulate some time-consuming operation
Elapsed time: 2.002082347869873 seconds
技术支持

在这里，Timer 类在进入 with 代码块时通过 __enter__ 方法记录开始时间。它返回 self，以便在代码块内访问该对象。退出 with 代码块时，__exit__ 方法计算并打印耗时。

基于函数的方式

如果您更偏好简洁的写法（可能以牺牲一定灵活性为代价），基于函数的方式是不错的选择。在这里，您可以使用 contextlib 模块中的 contextmanager 装饰器将任意函数转换为上下文管理器。

让我们看看它如何工作：

import time
from contextlib import contextmanager


@contextmanager
def timer():
   start_time = time.time()
   yield
   end_time = time.time()
   elapsed_time = end_time - start_time
   print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with timer():
       time.sleep(2)
Elapsed time: 2.0020740032196045 seconds
技术支持

@contextmanager 装饰器将 timer 函数转化为上下文管理器。在函数内部，先记录 start_time，随后 yield 语句会暂停执行，让 with 代码块中的代码运行。

最后，通过记录结束时间并打印耗时来实现 __exit__ 的功能。

本质上，您在 yield 之前编写 __enter__ 的逻辑，而在 yield 之后编写 __exit__ 的逻辑。这两种方式都能实现相同的目标，选择取决于您对结构与可读性的偏好。

上下文管理器的实用示例

既然您已经了解了如何在 Python 中编写自定义上下文管理器，让我们看看一些它们非常有用的实践场景。

管理文件操作

文件操作在许多应用中都很常见。无论是读写文件，恰当的管理都能确保数据完整性和资源效率。考虑这样一个场景：我们希望使用自定义上下文管理器向文件写入文本。

class FileManager:
   def __init__(self, filename, mode):
       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_value, traceback):
       if self.file:
           self.file.close()


# Example usage
if __name__ == "__main__":
   with FileManager("example.txt", "w") as file:
       file.write("Hello, world!\n")
技术支持

在这个示例中，我们定义了一个 FileManager 类来模拟 Python 中处理文件操作的 open 函数。下面按方法拆解其作用：

• __init__(self, filename, mode)：类的构造方法。用给定的文件名和模式初始化 FileManager 对象，指定要打开的文件以及打开模式（如读取、写入、追加）。
• __enter__(self)：如常，在进入 with 代码块时调用。按 mode 指定的模式打开 filename 对应的文件。打开的文件对象赋给 self.file 并返回，以便在 with 代码块中使用。
• __exit__(self, exc_type, exc_value, traceback)：在退出 with 代码块时调用，无论块内是否发生异常。它通过调用 self.file.close()（当 self.file 非 None 时）确保正确关闭文件，防止资源泄漏并确保清理到位。

执行上述代码后，如果您检查 example.txt 的内容，应能看到写入的文本：

$ cat example.txt
Hello, world!
技术支持

当然，FileManager 类的优势在于您可以为文件操作实现任意自定义逻辑。例如，您可以修改它，使其将打开或关闭了哪些文件记录到一个文本日志中。

import datetime


class FileManager:
   def __init__(self, filename, mode, log_filename="file_log.txt"):
       self.filename = filename
       self.mode = mode
       self.log_filename = log_filename
       self.file = None

   def log_action(self, action):
       timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
       with open(self.log_filename, "a") as log_file:
           log_file.write(f"{timestamp} - {action}: {self.filename}\n")

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

   def __exit__(self, exc_type, exc_value, traceback):
       if self.file:
           self.file.close()
           self.log_action("Closed")


# Example Usage
with FileManager("example.txt", "r") as file:
   content = file.read()
   print(content)

Hello, world!
技术支持

在这个修改版本中：

• 导入了 datetime 模块以处理时间戳。
• 新增 log_action 方法用于记录文件操作。
• 新增 log_filename 参数用于指定日志文件名，默认值为 'file_log.txt'。
• 在 __enter__ 和 __exit__ 方法中调用 log_action，将文件名和动作（打开或关闭）追加到日志文件。
• 在 log_action 方法中，使用 datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") 获取当前时间戳，格式为 "YYYY-MM-DD HH:MM:SS"。
• 随后将时间戳与文件名和动作一并写入日志。

现在，如果我们打印 file_log.txt 的内容，应能看到日志：

$ cat file_log.txt
2024-04-03 12:37:02 - Opened: example.txt
2024-04-03 12:37:02 - Closed: example.txt
技术支持

我们还可以将另一个实用场景实现为自定义上下文管理器：一个 SQLite 数据库管理工具：

import sqlite3


class DatabaseConnection:
   def __init__(self, database_name):
       self.database_name = database_name
       self.connection = None

   def __enter__(self):
       self.connection = sqlite3.connect(self.database_name)
       return self.connection

   def __exit__(self, exc_type, exc_value, traceback):
       if self.connection:
           self.connection.close()
技术支持

这同样是在模拟 sqlite3.connect 函数，但可以通过添加自定义的数据库管理逻辑进一步完善。让我们用一些示例函数来检查该管理器是否正常工作：

# Example usage
def create_table():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           """CREATE TABLE IF NOT EXISTS users (
                           id INTEGER PRIMARY KEY,
                           username TEXT,
                           email TEXT)"""
       )


def insert_data(username, email):
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           "INSERT INTO users (username, email) VALUES (?, ?)", (username, email)
       )
       connection.commit()


def fetch_data():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute("SELECT * FROM users")
       return cursor.fetchall()


if __name__ == "__main__":
   create_table()
   insert_data("john_doe", "[email protected]")
   insert_data("jane_doe", "[email protected]")
   users = fetch_data()
   print("Users in the database:")
   for user in users:
       print(user)

Users in the database:
(1, 'john_doe', '[email protected]')
(2, 'jane_doe', '[email protected]')
(3, 'john_doe', '[email protected]')
(4, 'jane_doe', '[email protected]')
技术支持

通过将数据库操作拆分为函数并使用我们刚刚创建的自定义管理器，我们可以编写出更易维护的代码。

与 Python 上下文管理器相关的高级主题

本节将讨论与 Python 上下文管理器相关的高级主题，包括错误处理与嵌套上下文管理器。

错误处理

错误处理是编程中不可或缺的一部分，而 Python 的上下文管理器为在资源管理操作中处理错误提供了便捷方法。

在 FileManager 类的示例中，可以在 __exit__ 方法中实现错误处理，即使发生异常也能确保正确清理。比如，如果在向文件写入数据时（处于 with 代码块内）发生错误，我们仍希望确保文件被关闭以防止资源泄漏。我们可以这样修改 __exit__ 方法：

def __exit__(self, exc_type, exc_value, traceback):
   if self.file:
       self.file.close()
       if exc_type is not None:
           print(f"An error occurred: {exc_value}")
技术支持

这里，我们首先调用 self.file.close() 关闭文件以确保正确清理。然后检查是否发生了异常 (exc_type is not None)。如果发生异常，就打印错误信息以反馈问题所在。

其中：

• exc_type 是异常的类，例如 AttributeError 和 IndexError。
• exc_value 是带有消息的异常实例。
• traceback 是错误的完整回溯信息。

让我们用一段故意出错的代码来运行 FileManager 类：

with FileManager("example.txt", "r") as file:
   # Simulate a sample error
   content = file.read()
   print(content[150])

An error occurred: string index out of range
技术支持

因此，在自定义管理器中处理异常、实现自定义逻辑时，您可以参考如下示例工作流：

def __exit__(self, exc_type, exc_value, traceback):
   # Perform cleanup
   ...

   # Handle the exception
   if isinstance(exc_value, AttributeError):
       # Handle AttributeError here
       print(f"Exception message: {exc_value}")
   # Handle another exception
   elif isinstance(exc_value, AnotherError):
       ...
技术支持

嵌套上下文管理器

嵌套上下文管理器允许您在单个 with 语句中同时管理多个资源，每个资源都有各自的上下文管理器。这在处理涉及多种资源的复杂操作时尤为有用。让我们在示例中看看如何嵌套上下文管理器。

设想这样一个场景：我们想在同一上下文中打开文件并进行数据库操作。也许我们希望将某个 CSV 文件的内容作为行写入数据库。我们可以在一个 with 语句中嵌套使用 FileManager 和 DatabaseConnection 上下文管理器来实现：

if __name__ == "__main__":
   with FileManager("example.txt", "w") as file, DatabaseConnection("example.db") as connection:
       ...
技术支持

在这个示例中，FileManager 和 DatabaseConnection 上下文管理器被嵌套在同一个 with 语句中。这样我们就能同时打开文件并建立数据库连接。随后与文件写入和数据库有关的操作均在这一合并的上下文中完成。

需要注意的一点是代码的可读性。建议同时管理的上下文数量不超过两个。

总结

Python 中的上下文管理器在 with 语句内为资源管理提供了结构化的方法，确保资源被恰当分配与释放。

在本教程中，我们学习了上下文管理器的基础知识，走过了实践示例，并讨论了错误处理和嵌套等高级主题。通过掌握自定义上下文管理器的编写，您可以写出更简洁、更可靠的代码，改进错误处理，并高效管理资源。

编写自定义上下文管理器这一主题属于“高级 Python”范畴。以下是一些相关资源，帮助您进一步提升技能：

• Python 函数编写课程
• 成为 Python 程序员技能路径