Python进阶系列 - 20讲 with ... as:

44 阅读3分钟

上下文管理器是一个很好的资源管理工具。 它们允许您在需要时精确地分配和释放资源。 一个大家熟知的例子是 with open() 语句:

测试代码:

with open('notes.txt', 'w') as f:
    f.write('some todo...')

这将打开一个文件,并确保在程序执行离开 with 语句的上下文后自动关闭它。 它还处理异常并确保即使在出现异常的情况下也能正确关闭文件。 在内部,上面的代码翻译成这样的:

f = open('notes.txt', 'w')
try:
    f.write('some todo...')
finally:
    f.close()

我们可以看到使用上下文管理器和 with 语句更短更简洁。

上下文管理器的例子

  • 打开和关闭文件
  • 打开和关闭数据库连接
  • 获取和释放锁
from threading import Lock
lock = Lock()

# 传统的方式:
lock.acquire()
lock.release()

# 好的方式:
with lock:
    # do stuff

自定义上下文管理器

为了支持我们自己的类的 with 语句,我们必须实现 __enter____exit__ 方法。 当执行进入 with 语句的上下文时,Python 会调用 __enter__。 在这里应该获取并返回资源。当执行再次离开上下文时,__exit__ 被调用并释放资源。

代码:

class ManagedFile:
    def __init__(self, filename):
        print("初始化类", filename)
        self.filename = filename

    def __enter__(self):
        print(">>>进入上下文管理<<<")
        self.file = open(self.filename, "w")
        return self.file

    def __exit__(self, exc_type, exc_value, exc_traceback):
        if self.file:
            self.file.close()
        print("<<<退出上下文管理>>>")


with ManagedFile("notes.txt") as f:
    print("一些操作...")
    f.write("写入内容...")

结果:

初始化类 notes.txt
>>>进入上下文管理<<<
一些操作...
<<<退出上下文管理>>>

异常处理

如果发生异常,Python 会将类型、值和回溯传递给 exit 方法。 它可以在这里处理异常。 如果 __exit__ 方法返回 True 以外的任何内容,则 with 语句会引发异常。

代码:

class ManagedFile:
    def __init__(self, filename):
        print('开始', filename)
        self.filename = filename
    def __enter__(self):
        print('>>>进入上下文管理<<<')
        self.file = open(self.filename, 'w')
        return self.file
    def __exit__(self, exc_type, exc_value, exc_traceback):
        if self.file:
            self.file.close()
        print('执行:', exc_type, exc_value)
        print('<<<退出上下文管理>>>')

#没有异常
with ManagedFile('notes.txt') as f:
    print('一些操作...')
    f.write('写入内容...')
print('继续...')

# 跑出异常,文件仍然关闭。
with ManagedFile('notes2.txt') as f:
    print('一些操作...')
    f.write('写入内容...')
    f.do_something() # ❌这个方法会抛出异常。
print('继续...')

结果:

开始 notes.txt
>>>进入上下文管理<<<
一些操作...
执行: None None
<<<退出上下文管理>>>
继续...
开始 notes2.txt
>>>进入上下文管理<<<
一些操作...
执行: <class 'AttributeError'> '_io.TextIOWrapper' object has no attribute 'do_something'
<<<退出上下文管理>>>
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/var/folders/p9/ny6y7zmj0qdblv697s3_vj040000gn/T/ipykernel_68956/2387064930.py in <module>
     23     print('一些操作...')
     24     f.write('写入内容...')
---> 25     f.do_something() # ❌这个方法会抛出异常。
     26 print('继续...')

AttributeError: '_io.TextIOWrapper' object has no attribute 'do_something'

我们可以在 __exit__ 方法中处理异常并返回 True

代码:

class ManagedFile:
    def __init__(self, filename):
        print('开始', filename)
        self.filename = filename
    def __enter__(self):
        print('>>>进入上下文管理<<<')
        self.file = open(self.filename, 'w')
        return self.file
    def __exit__(self, exc_type, exc_value, exc_traceback):
        if self.file:
            self.file.close()
        if exc_type is not None:
            print('管理器内:处理异常')
        print('<<<退出上下文管理>>>')
        return True
with ManagedFile('notes2.txt') as f:
    print('一些操作...')
    f.write('写入内容...')
    f.do_something()
print('继续...')

结果:

开始 notes2.txt
>>>进入上下文管理<<<
一些操作...
管理器内:处理异常
<<<退出上下文管理>>>
继续...

实现上下文管理为生成器

除了编写一个类,我们还可以编写一个生成器函数并使用 contextlib.contextmanager 装饰器来装饰它。 然后我们也可以使用 with 语句调用该函数。 对于这种方法,函数必须在 try 语句中yield 资源,并且释放资源的 __exit__ 方法的所有内容现在都放在相应的 finally 语句中。

示列代码:

from contextlib import contextmanager
@contextmanager
def open_managed_file(filename):
    f = open(filename, 'w')
    try:
        print('>>>进入上下文管理<<<')
        yield f
    finally:
        print('<<<退出上下文管理>>>')
        f.close()

with open_managed_file('notes.txt') as f:
    print("一些操作...")
    f.write('写入内容...')

生成器首先获取资源。 然后它会暂时挂起自己的执行并 yields 资源,以便调用者可以使用它。 当调用者离开 with 上下文时,生成器继续执行并释放 finally 语句中的资源。