一文读懂Python上下文管理器

写在前面

Python 的上下文管理器用于管理资源。它允许您以一种高效的方式处理资源，使您的代码更易读，减少错误。例如，文件处理、网络连接和数据库连接。它提供了一种在资源不再需要时自动分配和释放资源的方式，确保正确清理和防止资源泄漏。

上下文管理器通常与 with 语句一起使用，该语句确保由上下文管理器管理的资源在块的结束时被正确释放，即使发生异常也是如此。

Python 提供了两种主要方式来创建上下文管理器：通过定义一个带有__enter__和__exit__方法的类；使用 contextlib 模块的 contextmanager 装饰器。

自定义上下文管理器

让我们看下面的例子 –

f = open('sample.txt', 'w')
f.write('This is test.')
f.close()

在这个例子中，我们必须在使用文件之后手动关闭文件。

而使用上下文管理器，即使出现错误，也会自动关闭文件。

with open('sample.txt', 'w') as f:
    f.write('This is test.')

这个方法不仅对文件有用，还可以用于以下情况：

连接和关闭数据库连接
获取或释放锁
文件处理
网络连接

此外，我们可以编写自己的上下文管理器来处理自定义资源。我们可以使用类或函数和装饰器。

使用类

class OpenFile:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

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

    def __exit__(self, exc_type, exc_val, traceback):
        self.file.close()


with OpenFile('sample.txt', 'w') as f:
    f.write('Testing')

print(f.closed)

当我们使用上下文管理器 with 调用该类时，它会进入__init__方法并设置属性。然后它运行__enter__方法中的代码，即打开文件并返回变量 file。所以，上下文管理器中的变量f设置了返回的变量，我们可以以任何我们喜欢的方式对其进行操作。现在，当我们退出该块时，它调用__exit__方法，即调用 self.file.close()。所以当您打印 f.closed 时，它将返回 True。

使用函数和装饰器

from contextlib import contextmanager

@contextmanager
def open_file(file, mode):
    f = open(file, mode)
    yield f
    f.close()

with open_file('sample.txt', 'w') as f:
    f.write('This is test.')

print(f.closed)

在这里，您可以看到函数open_file做了和类中相同的事情。我们在这里使用yield语句。yield 之前的所有内容与类中的__enter__方法相同。yield 是 with 语句中的代码将运行的代码，即它与类中的__enter__方法中返回文件相同。yield 语句之后的所有内容与类中的__exit__方法相同。我们正在做与类代码中相同的事情，即使用with语句并调用生成器函数。

其他使用场景

import os

cwd = os.getcwd()  # 获取当前工作目录
os.chdir('directory_1')  # 切换目录
print(os.listdir())  # 列出文件
os.chdir(cwd)  # 切换回当前目录

cwd = os.getcwd()
os.chdir('directory_2')
print(os.listdir())
os.chdir(cwd)

如您所见，我获取当前工作目录，然后切换到其他目录并执行一些操作，然后再切换回当前目录。这样保存目录、切换目录，然后再切换回原始目录是有点不方便的。每次都要记住这样做。

这很麻烦，我们可以使用自定义的上下文管理器。现在您不必担心获取当前目录、切换到另一个目录，然后在执行相应操作后切换回原始目录。

import os
from contextlib import contextmanager

@contextmanager
def change_dir(destination):
    try:
        cwd = os.getcwd()
        os.chdir(destination)
        yield
    finally:
        os.chdir(cwd)

with change_dir('directory_1'):
    print(os.listdir())
with change_dir('directory_2'):
    print(os.listdir())