【Python 元编程】自定义上下文管理器 ✨

前言📜

with 语句想必大家都不陌生，在使用 Python 进行文件操作时候，都少不了它。

如以下：

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

但是我想大部分人只知其然，不知其所以然。

在 Python 编程中，我们常常需要处理资源的获取和释放，如文件操作 📄、数据库连接 📊、网络请求 🌐 等。为了确保这些资源的正确管理和释放，Python 提供了一种强大的工具——上下文管理器（context manager） 🧹。

本篇文章将介绍上下文管理器的概念，以及如何自定义上下文管理器来提高代码的可维护性和可读性。

上下文管理器🛠️

什么是上下文管理器？🤔

上下文管理器是一种用于为进入和退出运行时上下文提供额外行为的对象 🧳。它通常用于处理资源分配和释放，以确保资源被正确管理，不会出现资源泄漏等问题 🚫。

上下文管理器的作用🛠️

自定义上下文管理器的要义在于提高代码的可维护性和可读性 📚，同时确保资源的正确管理和异常处理 🚑。它使得代码更加健壮 💪，减少了资源泄漏和错误的风险 🌪️。

通过使用上下文管理器，我们可以优雅地处理资源的获取和释放，使代码更加清晰和可靠 🌟。

上下文管理器的应用场景🏆

上下文管理器广泛应用于处理各种资源，包括但不限于：

• 文件操作 📂：打开文件后自动关闭，确保文件资源被释放。
• 数据库连接 🗃️：管理数据库连接的获取和关闭，防止连接泄漏。
• 网络请求 🌐：处理网络请求的异常情况，确保资源的正常释放。

自定义上下文管理器可以根据不同场景的需求进行定制，提供更灵活的资源管理方式 🤝。

为什么使用上下文管理器？🔍

我用一个故事来阐述为什么我们需要使用上下文管理器：

📖 故事开始：日常的挑战

小明是一个 Python 开发者，负责维护一个相对复杂的系统。在这个系统中，有一部分功能涉及到频繁地与数据库交互。

起初，小明的数据库操作代码看起来是这样的：

import sqlite3

# 打开数据库，执行查询，然后关闭数据库
def query_database():
    db = sqlite3.connect('my_database.db')
    cursor = db.cursor()
    cursor.execute('SELECT * FROM my_table')
    result = cursor.fetchall()
    db.close()
    return result

这段代码虽然简单直接，但随着系统的发展，小明发现每次操作数据库都需要重复这个“打开-操作-关闭”的过程。

🚧 面临的问题

小明开始察觉到，随着代码的增长，这种模式变得越来越繁琐。

特别是，小明需要在每次数据库操作中包含错误处理来确保数据库连接的正确关闭：

import sqlite3

def query_database():
    """查询数据"""
    db = sqlite3.connect('my_database.db')
    try:
        cursor = db.cursor()
        cursor.execute('SELECT * FROM my_table')
        return cursor.fetchall()
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        db.close()

        
def insert_into_database(query, values):
    """插入数据"""
    db = sqlite3.connect('my_database.db')
    try:
        cursor = db.cursor()
        cursor.execute(query, values)
        db.commit()
    finally:
        db.close()


def update_database(quert, values):
    """更新数据"""
	# 省略

使用 try 和 except 来捕捉所有的异常看起来是个好办法，但是随着更多类似的操作被添加到项目中，小明发现自己在重复编写大量类似的错误处理代码。

💡 灵光一闪：上下文管理器

于是小明开始寻求优化方法，他了解到了 Python 中的一个特性——上下文管理器。发现它可以自动管理资源，无论过程中发生什么情况，都能保证资源（如文件、网络连接或数据库连接）被适当地清理。

🌟 改进方案：自定义上下文管理器

PS：原理在文章后面再做解释。

激动之下，小明决定为数据库连接创建一个自定义的上下文管理器。这个上下文管理器将会负责打开数据库连接、在操作完成后自动关闭连接，甚至在发生异常时确保安全地关闭连接。

注：这份自定义上下文管理器与 使用 try 和 except 来捕捉所有的异常 的代码效果一致，区别是更加优雅。

import sqlite3

class DatabaseConnection:
    def __enter__(self, db_path):
        self.connection = sqlite3.connect(db_path)	# 'my_database.db'
        return self.connection.cursor()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.connection.close()

现在，小明可以使用这个上下文管理器来简化数据库操作：

• 如查询和插入数据

with DatabaseConnection() as cursor:
    cursor.execute('SELECT * FROM my_table')
    result = cursor.fetchall()

with DatabaseConnection() as cursor:
    cursor.execute('INSERT INTO my_table VALUES (?)', ('value',))
    result = cursor.fetchall()

✨ 效果显著：清晰与安全

这种方法不仅使代码变得更加简洁，而且更加安全可靠。小明不再需要担心数据库连接是否正确关闭，甚至在遇到错误时也能保证资源的正确释放。

🏁 结论：优雅的资源管理

通过这个故事，我们看到上下文管理器如何帮助简化资源管理的代码，并增加代码的健壮性。它们提供了一种优雅和有效的方式来处理资源，确保即使在出现错误的情况下也能正确地管理资源。

这样，小明的系统变得更加健壮，代码也更加优雅和可维护。上下文管理器真的是 Python 编程中的一颗宝石！🚀👩‍💻👨‍💻

代码实现🔧

而实现上下文管理器，通常由两种方式：使用 class 类 或 使用 contextlib 模块。

注意事项

1. 一个类要成为上下文管理器，则必须实现两个特殊方法的对象构成：__enter__ 和 __exit__ 🎉。

   • __enter__ 方法在进入上下文时被调用，通常是资源的分配或初始化，然后返回一个值，该值将被传递给 as 关键字后的变量。

   • __exit__ 方法在退出上下文时被调用，通常用于资源的释放或异常处理。接受三个参数：异常类型、异常值和异常追踪信息。

     • exc_type（异常类型）：这个参数用于接收在上下文管理器内部代码块中引发的异常的类型（Exception 的子类）；
     • exc_value（异常值）：这个参数用于接收异常的实例，即异常的具体对象；
     • traceback（回溯信息）：这个参数用于接收与异常相关的回溯信息，它是一个包含堆栈跟踪信息的对象；
     • 如果代码块内没有异常发生，这三个参数的值将为 None。

2. 返回值：__enter__ 方法可以返回一个值，该值将被赋给 as 关键字后的变量。如果不需要返回值，可以返回 self 或者 None。

3. 异常处理：在 __exit__ 方法中，要根据需要处理异常。

   • 如果 __exit__ 方法返回 True，则表示异常已被处理，
   • 如果 __exit__ 方法返回 False，否则异常将被传播；
   • 通常，在 __exit__ 方法中，我们可以检查异常类型并根据需要进行处理，然后返回 True 或 False。

使用 Class 类实现🏗️

基础版

class MyContextManager:
    def __enter__(self):
        print("上下文管理器 ==> 初始化")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print("上下文管理器 ==> 退出啦")
        return True

调用这个上下文管理器，

with MyContextManager() as manager:
    print("上下文管理器 ==> 在内部")
    int('abc')

输出：

• 可以看到，执行的步骤~
• 并且， int(abc) 的异常也正常处理，程序并没有抛出异常~

上下文管理器 ==> 初始化
上下文管理器 ==> 在内部
上下文管理器 ==> 退出啦

但如果我们修改一下，

def __exit__(self, exc_type, exc_value, traceback):
    print("上下文管理器 ==> 退出啦")
    return False

那就可以看到异常被传播了

上下文管理器 ==> 初始化
上下文管理器 ==> 在内部
上下文管理器 ==> 退出啦
Traceback (most recent call last):
  File "F:\csdn\demo.py", line 13, in <module>
    int('abc')
ValueError: invalid literal for int() with base 10: 'abc'

复杂版

在处理数据库事务时，确保事务要么完全执行，要么在出现错误时完全回滚，同时确保数据库连接被正确关闭。

# encoding=utf-8

import pymysql


# 自定义上下文管理器，用于MySQL数据库连接
class MySQLDB:
    def __init__(self, host='localhost', user='root', password='123456', database='demo'):
        self.host = host
        self.user = user
        self.password = password
        self.database = database

    def __enter__(self):
        self.connection = pymysql.connect(
            host=self.host,
            user=self.user,
            password=self.password,
            database=self.database
        )
        self.cursor = self.connection.cursor()
        return self.cursor

    def __exit__(self, exc_type, exc_value, traceback):
        self.cursor.close()
        self.connection.close()
        if exc_type is not None:
            print(f"Exception inside the context: {exc_type}, {exc_value}")


# 不使用上下文管理器，手动管理MySQL数据库连接
def manually_manage_mysql():
    connection = pymysql.connect(
        host='localhost',
        user='root',
        password='123456',
        database='demo'
    )
    cursor = connection.cursor()
    try:
        cursor.execute("SELECT * FROM test")
        result = cursor.fetchall()
        # 触发一个异常
        int("invalid")
    except Exception as e:
        print(f"Exception inside the context: {e}")
    finally:
        cursor.close()
        connection.close()


if __name__ == '__main__':
    # 使用上下文管理器自动管理MySQL数据库连接
    with MySQLDB() as cursor:
        cursor.execute("SELECT * FROM test")
        results = cursor.fetchall()
        # 触发一个异常
        int("invalid")

    print("\n" + "=" * 40 + "\n")
    # 不使用上下文管理器，手动管理MySQL数据库连接
    manually_manage_mysql()

代码释义：

1. 上下文管理器的创建：在这里，创建了一个名为MySQLDB的自定义上下文管理器类。它包含__enter__和__exit__方法，分别用于在进入和离开上下文时执行操作。

2. 资源管理：在__enter__方法中，建立了与MySQL数据库的连接，并返回一个光标对象。这个光标对象用于执行SQL查询。

3. 异常处理：在__exit__方法中，关闭了光标和数据库连接。如果在上下文中出现异常，__exit__方法会捕获并处理这些异常，并确保资源的正确释放。在示例中，我们使用exc_type、exc_value和traceback参数来捕获异常信息，并在异常发生时打印出来。

4. 使用上下文管理器：通过使用with语句，我们可以轻松地管理MySQL数据库连接。在with块内，资源会在进入和离开上下文时自动管理，无需手动处理连接的打开和关闭。

5. 手动管理资源：还展示了不使用上下文管理器的情况下如何手动管理MySQL数据库连接。在这个示例中，需要手动打开和关闭连接，以及处理异常情况。这种方式更容易出错和繁琐。

优势：

上面的两份代码作用几乎一致，但显然可以看到，使用上下文管理器的优点：

1. 可读性和简洁性：使用上下文管理器的代码更加简洁和易于理解。上下文管理器在进入和离开上下文时自动处理连接的创建和关闭，无需显式调用cursor.close()和connection.close()，使代码更加清晰。

2. 异常处理：上下文管理器能够处理异常情况，确保在发生异常时关闭连接。在第一个代码示例中，如果在上下文中发生异常，__exit__方法会自动关闭游标和连接，同时记录异常信息。而在第二个手动管理连接的示例中，需要手动编写异常处理代码，容易忽略或出现错误。

3. 代码重用性：使用自定义上下文管理器的方式可以更容易地在多个地方重用数据库连接的管理逻辑。而手动管理连接的方式需要在每次使用时复制和粘贴连接代码。

4. 规范性：上下文管理器遵循了 Python 的上下文管理协议，符合 Python 的编程规范，使代码更加规范和可维护。

总的来说，自定义上下文管理器的要义在于提高代码的可维护性和可读性，同时确保资源的正确管理和异常处理。它使得代码更加健壮，减少了资源泄漏和错误的风险。在实际应用中，上下文管理器可用于处理各种资源，包括文件、数据库连接、网络连接等。

使用 contextlib 模块🏹

这份代码等价于 上面 🌟 改进方案：自定义上下文管理器 的代码。

使用contextlib模块中的contextmanager装饰器创建上下文管理器时，有几个关键点需要注意：

1. 生成器函数：
   • 使用contextmanager装饰器的函数应该是一个生成器函数，它通过yield语句产生一个值，该值会被with语句的目标变量接收。
2. 资源管理：
   • 在yield之前的代码块通常用于设置或获取资源（例如打开文件、建立数据库连接等）。
   • 在yield之后的代码块，即finally部分，用于清理或释放资源（例如关闭文件、断开数据库连接等）。
3. 异常处理：
   • 在yield语句之后的代码将在退出with块时执行，无论with块中是否发生异常。你可以在这部分代码中添加异常处理逻辑。
   • 如果with块中发生异常，它会被重新抛出，除非在退出部分（finally块）明确捕获并处理。
4. 只能产生一个值：
   • 用contextmanager装饰的函数应该只包含一个yield语句。在yield之后的代码会在with块完成时执行。

！！！注意在函数前面加上@contextmanager

from contextlib import contextmanager

import sqlite3


@contextmanager
def database_connection(db_path):
    connection = sqlite3.connect(db_path)
    try:
        yield connection.cursor()
    finally:
        connection.close()

🌈 总结

通过本文，我们深入了解了 Python 中上下文管理器的强大功能和实际应用。上下文管理器不仅提高了代码的可读性和可维护性，还确保了资源的正确管理，避免了资源泄漏。它是 Python 编程中一个重要的工具，尤其在处理需要精细管理资源的场景中显示出其价值。

自定义上下文管理器使得代码更加简洁、清晰，也更加健壮。通过使用 class 或 contextlib 模块，我们可以根据不同的需求创建灵活且强大的上下文管理器，提高代码的整体质量和效率。

无论是在大型软件开发中，还是在日常的脚本编写中，合理使用上下文管理器都能带来显著的好处。掌握它们的使用，是每个 Python 开发者提升编程技能的重要一步。