本书作者：测试开发刚哥，文章版权归个人所有，未经授权请勿转载。

with 语句会设置一个临时的上下文，交给上下文管理器对象控制，并且负责清理上下文。

比如在打开文件时通常都会使用 with 语句：

with open("a.txt") as f:
    f.read()

with 块执行完后会自动关闭文件。

Python 上下文管理器对象存在的目的就是管理 with 语句。

实现一个上下文管理器

上下文管理器协议包含__enter__和__exit__两个方法，所以要实现一个上下文管理器，就得实现这两个方法，比如：

class LookingGlass:

    def __enter__(self):
        import sys
        self.original_write = sys.stdout.write
        sys.stdout.write = self.reverse_write
        return 'JABBERWOCKY'

    def reverse_write(self, text):
        self.original_write(text[::-1])

    def __exit__(self, exc_type, exc_value, traceback):
        import sys
        sys.stdout.write = self.original_write
        if exc_type is ZeroDivisionError:
            print('Please DO NOT divide by zero!')
            return True
  • __enter__是上下文管理器的入口，在 with 语句开始运行时调用。
  • __exit__是上下文管理器的出口，在 with 语句运行结束后调用。

上下文管理器对象是在执行 with 后面的表达式得到的：

with LookingGlass() as what:
    print(what)

as 子句是可选的，在 as 执行前就已经得到了上下文管理器对象，在 as 执行时实际就已经调用__enter__了。

把生成器变为上下文管理器

使用@contextmanager装饰器能减少创建上下文管理器的样板代码量，只需要实现一个有 yield 语句的生成器，生成想让__enter__方法返回的值。

示例：

import contextlib


@contextlib.contextmanager
def looking_glass():
    import sys
    original_write = sys.stdout.write

    def reverse_write(text):
        original_write(text[::-1])

    sys.stdout.write = reverse_write
    msg = ''
    try:
        yield 'JABBERWOCKY'
    except ZeroDivisionError:
        msg = 'Please DO NOT divide by zero!'
    finally:
        sys.stdout.write = original_write
        if msg:
            print(msg)

yield 语句的作用是：

  • 在 yield 语句前面的代码在 with 块开始时执行，相当于__enter__。
  • 在 yield 语句后面的代码在 with 块结束时执行，相当于__exit__。

注意这里的 yield 与迭代没有任何关系，这其实引出了 Python 另一个重要技术点，协程：执行到某一点时暂停，让客户代码运行，直到客户让协程继续做事。下篇文章将展开对 Python 协程的学习。

参考资料：

《流畅的 Python》第 15 章 上下文管理器和 else 块