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

您可能已经熟悉with语句，这是一种在Python中处理资源的简洁方式。但你有没有想过它是如何工作的？with关键字的强大功能来自上下文管理器。

上下文管理器是Python中的一个基本设计模式，它提供了一种结构化的资源管理方法。它们确保资源被获取、正确使用，然后最终释放或清理，即使有错误或异常。

在处理文件、网络连接或数据库句柄等资源时，这一点尤为重要。通过使用上下文管理器，我们可以编写更干净、更可靠的代码，从而摆脱忘记关闭文件或释放锁的担忧。

在本文中，我们将超越Python提供的默认上下文管理器，并学习编写自定义上下文管理器。

理解Python中的上下文管理器

在底层，上下文管理器是定义了两个特殊方法的对象：__enter__和__exit__。当你进入with块时，__enter__方法被调用，它的返回值被赋给该块中的一个变量。另一方面，__exit__方法在with块退出时被调用，不管它是正常完成还是异常完成。

这种结构确保了正确的资源处理。让我们来看看Python中的一些内置上下文管理器，以更好地理解这一点：

1.文件管理

举一个打开文件的经典例子：

with open('file.txt', 'r') as file:
   data = file.read()

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

接下来，file对象的__exit__方法保证在with块退出时被调用，即使发生异常。此方法负责关闭文件，确保您不会留下打开的文件句柄。

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__方法获取锁，确保一次只有一个线程可以执行。最后，__exit__方法在退出with块时释放锁，允许其他线程继续。

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。然后可以使用游标与数据库交互。__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: 0.0635 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类定义了__enter__方法，以在您进入with块时捕获开始时间。它返回self以允许访问块内的对象。__exit__方法计算退出with块时所用的时间并打印出来。

基于函数的方法

如果您更喜欢更简洁的方法（可能以牺牲灵活性为代价），那么基于函数的方法可能是一个不错的选择。这里，您使用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: 0.0635 seconds")


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

@contextmanager装饰器将timer函数转换为上下文管理器。在函数内部，start_time被捕获，yield语句暂停执行，允许with块中的代码运行。

最后，__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不为None，则通过调用self.file.close()来确保文件正确关闭。这可以防止资源泄漏并确保正确的清理。

如果你在执行上述代码后检查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-18 21:37:02 - Opened: example.txt
2024-04-18 21: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", "john@example.com")
   insert_data("jane_doe", "jane@example.com")
   users = fetch_data()
   print("Users in the database:")
   for user in users:
       print(user)

Users in the database:
(1, 'john_doe', 'john@example.com')
(2, 'jane_doe', 'jane@example.com')
(3, 'john_doe', 'john@example.com')
(4, 'jane_doe', 'jane@example.com')

总结

Python中的上下文管理器在with语句中提供了一种结构化的资源管理方法，确保资源的正确分配和释放。

在本文中，我们学习了上下文管理器的基础知识。通过掌握编写自定义上下文管理器，您可以编写更干净、更可靠的代码，改进错误处理，并有效地管理资源。