你真的了解python的with语句吗？通过分析contextlib源码让你彻底掌握with的用法

楔子

下面我们来聊一下Python中的上下文管理，Python中的上下文管理我们可以通过with语句实现。在Python中使用with语句最多的情况，莫过于操作文件了，比如我们在打开一个文件的时候会通过类似于with open("test.txt", encoding="utf-8") as f: 这种形式打开，这种方式的好处就在于with语句结束后会自动关闭文件。

那么with语句的原理是什么呢？我们怎么样才能使用with语句呢？这次我们就全方位的剖析一下，并且在Python的标准库中还有一个模块叫做contextlib，从名字上也能看出来这是一个用于上下文管理的模块。我们后面也会通过分析contextlib的源码，来自己实现一下contextlib的功能。

上下文管理器API

上下文管理器(context manager)负责管理一个代码块中的资源，会在进入代码块时创建资源，然后再退出代码后清理这个资源。比如：文件就支持上下文管理器API，可以确保文件读写后关闭文件。

下面我们就来模拟一下文件的读取

class Open:

    def __init__(self, filename, mode='r', encoding=None):
        self.filename = filename
        self.mode = mode
        self.encoding = encoding

    def __enter__(self):
        print("__enter__, 有了这个就可以使用with Open() as f语句，这里的f就是我return的内容")
        return self

    def read(self):
        print(f"文件进行读操作,读取文件：{self.filename!r}, 模式：{self.mode!r}, 编码：{self.encoding!r}")

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("__exit__, 我是用来清理资源的，当操作执行完毕之后就会执行我，比如：关闭文件")


with Open("test.txt", encoding="utf-8") as f:
    f.read()
"""
__enter__, 有了这个就可以使用with Open() as f语句，这里的f就是我return的内容
文件进行读操作,读取文件：'test.txt', 模式：'r', 编码：'utf-8'
__exit__, 我是用来清理资源的，当操作执行完毕之后就会执行我，比如：关闭文件
"""

我们看到with Open() as f:的流程就是，先实例化一个Open对象，然后通过实例对象来调用__enter__方法，将其返回值赋值给with语句中的f，然后执行with语句块中的代码，最后执行__exit__方法。

但是注意的是：with Open() as f中的这个f是什么，取决于__enter__中返回了什么。

class Open:

    def __enter__(self):
        return "古明地觉"

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


# 如果是f = Open(), 那么毫无疑问这个f就是类Open的实例对象
# 但是对于with Open() as f, 我们说这个f到底是什么, 它取决于__enter__中返回了什么
with Open() as f:
    print(f)
"""
古明地觉
"""

# 我们看到print(f)打印的是一个字符串，因为__enter__中返回的就是一个字符串
# 首先with Open() as f:这一行代码所做的事情就是先实例化一个Open对象
# 只不过这个实例对象我们没有用变量进行接收, 但它确实存在
# 然后该实例对象再调用__enter__, 将__enter__的返回值赋值给f
# 因此__enter__返回了什么, 这个f就是什么, 所以在with代码块中打印f得到的是一个字符串
# 所以要记住: f是由__enter__的返回值决定的，只不过大多数情况下，__enter__里面返回的都是self本身，所以相应的f指向的也是该类的实例对象
# 然后with语句结束，也是通过这个实例对象来调用__exit__，这一点需要记清楚


# 当然，我们先实例化一个对象，再使用with也是可以的
o = Open()
with o as f:
    print(f"{f}, 世界第一可爱")
"""
古明地觉, 世界第一可爱
"""

# 如果是with Open() as f:, 那么实例化和调用__enter__就放在一起执行了
# 如果直接对一个实例对象使用with语句
# 比如with o as f:, 那么会直接通过实例对象o来调用__enter__, 将其返回值赋值给f
# 当with语句结束，再通过实例对象o来调用__exit__，进行资源的释放等操作


# 当然直接with o:, 不通过as f接收__enter__的返回值也是可以的
with o:
    pass

因此我们看到，一个对象究竟能否使用with语句，取决于实例化该对象的类(或者继承的基类)中是否同时实现了__enter__和__exit__两个魔法函数，两者缺一不可。

因此with语句的流程我们就很清晰了，以with XXX() as xx:为例，总共分为三步：

• 创建XXX()的实例对象，然后调用__enter__方法，将其返回值交给as xx中的xx
• 执行with语句块的代码
• 最后由该实例对象再调用__exit__进行一些收尾工作。

__enter__我们清楚了，但是我们发现__exit__里面除了self之外，还有三个参数分别是exc_type, exc_val, exc_tb，它们是做什么的呢？显然这三个参数分别是异常类型、异常值、异常的回溯栈， 从名字上也能看出来。

class Open:

    def __enter__(self):
        return "古明地觉"

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 这几个参数是使用pycharm的时候，自动帮我们加上去的
        print("__exit__执行:")
        print(exc_type)
        print(exc_val)
        print(exc_tb)
        return True


with Open() as f:
    print(f)
"""
古明地觉
__exit__执行:
None
None
None
"""

# 我们看到exc_type, exc_val, exc_tb三者全部为None
# 因为它们是和异常有关的，而我们这里没有出现异常，所以为None
# 但如果出现异常了呢？
with Open() as f:
    print(f)
    1 / 0
    print(123)
    print(456)
    print(789)
print("你猜我会被执行吗？")
"""
古明地觉
__exit__执行:
<class 'ZeroDivisionError'>
division by zero
<traceback object at 0x0000024CD4E4C080>
你猜我会被执行吗？
"""

我们看到在没有出现异常的时候，exc_type, exc_val, exc_tb打印的值全部是None。然而一旦with语句里面出现了异常，那么会立即执行__exit__，并将 异常的类型，异常的值，异常的回溯栈 传入到__exit__中。

因此：当with语句正常结束之后会调用__exit__，如果with语句里面出现了异常则会立即调用__exit__。

但是__exit__函数返回了个True是什么意思呢？当with语句里面出现了异常，理论上是会报错的，但是由于要执行__exit__函数，所以相当于暂时把异常塞进了嘴里。

如果__exit__函数最后返回了一个布尔类型为True的值，那么会把塞进嘴里的异常吞下去，程序不报错正常执行。如果返回布尔类型为False的值，会在执行完__exit__函数之后再把异常吐出来，引发程序崩溃。

这里我们返回了True，因此程序正常执行，最后一句话被打印了出来。但是 1 / 0 这句代码后面的几个print却没有打印，为什么呢？

因为我们说上下文管理执行是有顺序的，1) 先实例化Open的实例对象，调用__enter__方法，将__enter__的返回值交给f，2) 执行with语句块的代码，3) 最后调用__exit__

只要__exit__函数执行结束，那么这个with语句就算结束了。而with语句里面如果有异常，那么会立即进入__exit__函数，因此异常语句后面的代码是无论如何都不会被执行的。另外，如果__enter__中出现了异常，那么with语句会直接报错，同理__exit__中如果出现了异常也会报错，即使最后返回了True。

目前我们就把with说完了，下面我们进行contextlib的分析。我们说contextlib是一个专门用于上下文管理的内置模块，我们就来分析一下它内部是怎么实现的。

contextlib实现机制的分析

上下文管理器作为函数修饰符

contextlib中有一个类ContextDecorator，增加了对常规上下文管理器类的支持，使得上下文管理器，也可以作为函数的装饰器，我们来看一下。

import contextlib


class Context(contextlib.ContextDecorator):
    def __init__(self, how_used):
        self.how_used = how_used

    def __enter__(self):
        print(f"__enter__({self.how_used})")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"__exit__({self.how_used})")
        return True


@Context("我要去装饰了")
def foo(name):
    print(123)
    return f"我的名字叫: {name}"


print(foo("古明地觉"))
"""
__enter__(我要去装饰了)
123
__exit__(我要去装饰了)
我的名字叫: 古明地觉
"""

我们看到上下文管理器还可以作为函数的装饰器，我们看到先执行了__enter__，然后执行foo函数内部的代码，最后执行__exit__，打印返回值。

我们分析一下内部是如何实现的，首先我们装饰foo的时候，显然是使用Context的实例对象去装饰的，相当于给这个实例对象加上了括号，并且把foo这个函数作为参数传进去了。既然实例对象加上了括号(调用)，这就意味着该实例对象对应的类一定有__call__方法，但是我们定义的没有，那么继承的父类肯定有。

ContextDecorator这个类的代码很少，我们直接仿照一个。

from functools import wraps


class MyContextDecorator:

    def __call__(self, func):
        @wraps(func)
        def inner(*args, **kwargs):
            with self:
                return func(*args, **kwargs)
        return inner


class Context(MyContextDecorator):
    def __init__(self, how_used):
        self.how_used = how_used

    def __enter__(self):
        print(f"__enter__({self.how_used})")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"__exit__({self.how_used})")
        return True


@Context("我要去装饰了")
def foo(name):
    print(123)
    return f"我的名字叫: {name}"


print(foo("古明地觉"))
"""
__enter__(我要去装饰了)
123
__exit__(我要去装饰了)
我的名字叫: 古明地觉
"""

我们实现了一个类MyContextDecorator，实现了一模一样的效果。下面就来分析一下，整个流程。

class MyContextDecorator:

    def __call__(self, func):
        @wraps(func)
        def inner(*args, **kwargs):
            with self:
                return func(*args, **kwargs)
        return inner

类的源码很少，当我们使用实例对象去装饰foo函数的时候，就相当于给实例对象加上了括号

而类Context内部没有定义__call__, 那么肯定要走Context的父类也就是这里的__call__。

但是执行的时候，里面这个self可不是MyContextDecorator的实例对象，而是我们之前定义的Context类里面的self，也就是Context的实例对象。

如果面向对象基础不好的话，建议去复习一下，这里简单说一下。调用Context类实例对象的时候，肯定先去找Context里面的__call__方法，但是没有，那么会调用父类的，但是调用时对应的self还是Context的self

因此之前的@Context("我要作为装饰器去装饰")，就等价于context = Context("我要作为装饰器去装饰"); @context

然后再装饰foo的时候，相当于foo = context(foo)，那么会调用这里的__call__方法，然后foo会被传递给这里的func，然后返回inner。

所以foo在被装饰完之后，foo就变成了这里的inner，只不过由于有@wraps(func)这个装饰器在，所以装饰之后的函数名、__doc__等元信息没有改变

那么当我再调用foo("古明地觉")的时候，就等价于调用这里的inner("古明地觉")，而里面的with self:中的self，显然就是Context的实例对象。

所以就很清晰了，with self:，会先执行Context中的__enter__，然后执行func、也就是我们原本的foo，拿到它的返回值，然后执行__exit__，最后最外层的print再将拿到的返回值打印。

整体的逻辑就是上面分析的那样，希望能仔细理清一遍这里的流程。但是有一点需要注意：外层的print打印的是foo的返回值，不是__exit__的返回值。

@Context("我要去装饰了")
def foo(name):
    print(123)
    1 / 0
    return f"我的名字叫: {name}"


print(foo("古明地觉"))
"""
__enter__(我要去装饰了)
123
__exit__(我要去装饰了)
None
"""

我们看到foo中出现了异常，导致立刻执行了__exit__，而返回值则是一个None，因为foo执行失败了，所以结果就是一个None，不是__exit__里面的True。

从生成器到上下文管理器

采用传统方式创建上下文管理器并不难，只需要包含一个__enter__方法和一个__exit__方法的类即可。 不过某些时候，如果只有很少的上下文需要管理，那么完整地写出所以代码便会成为额外的负担。 在这些情况下，可以使用contextmanager修饰符将一个生成器函数转换为上下文管理器。

import contextlib


@contextlib.contextmanager
def foo(name, where):
    print(f"我的名字是: {name}, 居住在: {where}")
    yield "baka⑨"
    print(f"只要你喜欢{name}, 我们就是好兄弟")


with foo("古明地觉", "东方地灵殿") as f:
    print(f.upper())
"""
我的名字是: 古明地觉, 居住在: 东方地灵殿
BAKA⑨
只要你喜欢古明地觉, 我们就是好兄弟
"""

只要给函数加上这个装饰器，那么便可以使用with语句。当中的yield相当于将代码块分隔为两个战场：yield上面的代码相当于__enter__会先执行，然后将yield的值交给f，然后执行with语句，最后执行yield下面的代码块，相当于__exit__

注意：如果使用contextmanager装饰的话，函数中只能出现、且必须出现一个yield。

下面我们来手动实现一下contextmanager函数，contextlib中实现的比较复杂，主要是最后对异常进行了很多的检测。我们可以进行适当地简化一下，把主要的逻辑实现一下。

from functools import wraps


def my_contextmanager(func):

    class GeneratorContextManager:

        def __init__(self, func, *args, **kwargs):
            self.gen = func(*args, **kwargs)

        def __enter__(self):
            try:
                assert hasattr(self.gen, "__next__")
                return next(self.gen)
            except AssertionError:
                raise RuntimeError("函数中必须出现、且只能出现一个yield") from None

        def __exit__(self, exc_type, exc_val, exc_tb):
            try:
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("函数中必须出现、且只能出现一个yield")

    @wraps(func)
    def inner(*args, **kwargs):
        return GeneratorContextManager(func, *args, **kwargs)
    return inner


@my_contextmanager
def foo(name, where):
    print(f"我的名字是: {name}, 居住在: {where}")
    yield "baka⑨"
    print(f"只要你喜欢{name}, 我们就是好兄弟")


with foo("古明地觉", "东方地灵殿") as f:
    print(f.upper())
"""
我的名字是: 古明地觉, 居住在: 东方地灵殿
BAKA⑨
只要你喜欢古明地觉, 我们就是好兄弟
"""

我们手动的实现了一个contextmanager，下面还是分析一下整体的流程。

def my_contextmanager(func):

    class GeneratorContextManager:

        def __init__(self, func, *args, **kwargs):
            self.gen = func(*args, **kwargs)

        def __enter__(self):
            try:
                assert hasattr(self.gen, "__next__")
                return next(self.gen)
            except AssertionError:
                raise RuntimeError("函数中必须出现、且只能出现一个yield") from None

        def __exit__(self, exc_type, exc_val, exc_tb):
            try:
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("函数中必须出现、且只能出现一个yield")

    @wraps(func)
    def inner(*args, **kwargs):
        return GeneratorContextManager(func, *args, **kwargs)
    return inner

当使用my_contextmanager装饰的时候，外层的foo函数就变成了这里的inner。

然后with foo("古明地觉", "东方地灵殿") as f:的时候，等价于with inner("古明地觉", "东方地灵殿") as f:，等价于with GeneratorContextManager(foo, "古明地觉", "东方地灵殿") as f:。

然后创建一个GeneratorContextManager的实例对象，而self.gen显然是生成器函数foo对应的生成器，此时生成器已经创建。

我们创建完实例对象之后，要干啥来着，要执行__enter__。而我们说函数中必须出现yield，那么它是一个生成器函数，所以self.gen内部要有__next__方法，因此进行一个断言。

然后return next(value)开始驱动生成器运行了，当运行到yield的时候停下来。此时yield上面的代码已经执行完毕了，然后返回yield后面的值，赋值给f，开始执行with里面的代码。

with语句块的代码执行完毕，执行__exit__，在里面我们继续next，显然要驱动生成器继续执行，将yield下面的代码执行完毕，如果出现了StopIteration，那么什么也不做。如果没有出现此异常，证明不止一个yield，那么就抛出一个RuntimeError

所以我们看到，这个装饰器本质上还是使用了类的上下文管理。

由于contextmanager返回的上下文管理器派生自ContextDecorator，所以也可以被用作函数修饰符

@contextlib.contextmanager
def deco(name, where):
    print(f"我的名字是: {name}, 居住在: {where}")
    yield
    print(f"只要你喜欢{name}, 我们就是好兄弟")


@deco("古明地觉", "东方地灵殿")
def bar():
    print("猜猜我会在什么地方输出")

bar()
"""
我的名字是: 古明地觉, 居住在: 东方地灵殿
猜猜我会在什么地方输出
只要你喜欢古明地觉, 我们就是好兄弟
"""

当我执行bar的时候，还会先执行deco中yield上面的内容，然后执行bar代码的内容，最后执行deco中yield下面的内容。并且此时yield后面的内容是什么也已经无关紧要了，因为根本用不到了，yield不出去了。

那么我们如何实现上面的功能呢？可以当做练习自己尝试一下，并不难。

关闭打开的句柄

诸如打开文件之类的io操作，都会有一个close操作。因此为了确保关闭，可以使用contextlib中的一个叫做closing的类

class Open:

    def __init__(self):
        self.status = "open"

    def main(self):
        return "执行了很复杂的逻辑"

    def close(self):
        self.status = "closed"


with contextlib.closing(Open()) as f:
    print(f.main())  # 执行了很复杂的逻辑
    print(f"状态: {f.status}")  # 状态: open

# with语句结束后
print(f"状态: {f.status}")  # 状态: closed

contextlib.closing接收类的实例对象，其实主要就帮我们做了两件事：一个是可以通过with语句的方式来执行，另一个是执行完毕之后自动帮我们调用close方法。这里我们不手动实现了，非常简单，感觉没啥卵用，直接看源码中是如何实现的吧。

class closing(AbstractContextManager):

    def __init__(self, thing):
        # 这里的thing显然是我们之前传入的Open的实例对象f
        self.thing = thing
    def __enter__(self):
        # 先调用__enter__返回之前的实例
        return self.thing
    def __exit__(self, *exc_info):
        # 最后调用我们实例的close方法
        self.thing.close()

忽略异常

很多情况下，忽略产生的异常很有用，如果这个异常无法百分百避免、并且该异常又没啥卵用，那么该错误就可以被忽略。 要忽略异常，最常用的办法就是利用一个try except语句。但是在我们此刻的主题中，try except也可以被替换成contextlib.suppress()，以更显示地抑制with块中产生的异常

def foo():
    print(123)
    1 / 0
    print(456)


with contextlib.suppress(ZeroDivisionError, TypeError):
    foo()
    print(789)
"""
123
"""

在foo中出现了除零错误，但是程序并没有报错。相当于异常被"镇压"了，注意：如果with块中出现的异常，无法被suppress传入的异常捕获的话，那么异常还是会抛出来的。但是对于当前的例子来说，除零错误显然是被成功捕获了，最终只输出了123，可以看到不仅1/0下面的456没有被打印，连foo()下面的789也没有被打印。

因为只要出现了异常，就会进入到__exit__中，我们看一下源码是如何实现的。

class suppress(AbstractContextManager):
    def __init__(self, *exceptions):
        self._exceptions = exceptions

    def __enter__(self):
        pass

    def __exit__(self, exctype, excinst, exctb):
        return exctype is not None and issubclass(exctype, self._exceptions)

非常简单，没有发生异常就不说了，如果发生异常，但它是指定的某个异常的子类的话，程序也不会报错。注意：一个类也是其本身的子类，issubclass(TypeError, TypeError)结果为True。

异步上下文管理器的实现

Python在3.5的时候引入了async和await，可以通过async def定义一个原生的协程函数，通过await驱动一个协程执行。

而异步上下文则可以通过async with来实现

import asyncio


class A:

    def __init__(self, name):
        self.name = name

    async def __aenter__(self):
        print("__aenter__")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        print("__aexit__")
        return True


# 必须定义一个协程函数，然后事件循环驱动协程执行
async def main():
    async with A("古明地觉") as f:
        print("<<<>>>")


asyncio.run(main())
"""
__aenter__
<<<>>>
__aexit__
"""

当然contextlib这个包里面也实现与异步上下文管理相关的类和函数。

import contextlib
import asyncio


@contextlib.asynccontextmanager
async def foo():
    print(123)
    yield 456
    print(789)


async def main():
    async with foo() as f:
        print(f == 456)


asyncio.run(main())
"""
123
True
789
"""

关于异步上下文管理，其实和普通的同步上下文管理是类似的，有兴趣可以自己实现一下。