深度剖析下上下文管理器contextmanager

发表于 | 分类于 | | 阅读次数: |
字数统计: 3,075 字 | 阅读时长 ≈ 13 分钟

我们首先看下今天探讨的对象contextmanager的源码。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def contextmanager(func):
    """@contextmanager decorator.

    Typical usage:

        @contextmanager
        def some_generator(<arguments>):
            <setup>
            try:
                yield <value>
            finally:
                <cleanup>

    This makes this:

        with some_generator(<arguments>) as <variable>:
            <body>

    equivalent to this:

        <setup>
        try:
            <variable> = <value>
            <body>
        finally:
            <cleanup>
    """
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper

我们看到contextmanager可以将一个生成器函数修饰成上下文管理器。 这是如何实现的呢？

在剖析contextmanager之前我们要回顾两个知识点：上下文管理器、生成器。

何为上下文管理器？

对于实现了__enter__和__exit__这两个上下文管理器协议的对象就是一个上下文管理器。

1
2
3
4
5
6
7
8
9
10
11
class Contextor:
    def __enter__(self):
        pass
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

contextor = Contextor()

with contextor [as var]:
    with_body

上述代码的执行过程等价于：

1
2
3
4
5
6

VAR = contextor.__enter__()
try:
    with_body
finally:
    contextor.__exit__()

配合with语句使用的时候，上下文管理器会自动调用__enter__方法，然后进入运行时上下文环境，如果有as 从句，返回自身或另一个与运行时上下文相关的对象，值赋值给var。当with_body执行完毕退出with语句块或者with_body代码块出现异常，则会自动执行__exit__方法，并且会把对于的异常参数传递进来。如果__exit__函数返回True。则with语句代码块不会显示的抛出异常，终止程序，如果返回None或者False，异常会被主动raise，并终止程序。

我们看下__exit__参数含义：

exc_type是异常的类。 exc_val是异常实例。 exc_tb是一个追溯对象，在types.TraceBackType有一个引用。一般情况下:

  • type(exc_val) is exc_type
  • exc_val.__traceback__ is exc_tb

请注意，当上下文管理器下的代码未引发任何异常时，仍会调用__exit__ ，并且args将为(None, None, None)因此所有三个参数都应标注为optional 。

我们主要复习下 当with body出现异常的时候 异常将传递到 __exit__

何为生成器？

在明白contextmanager是如何将一个生成器变成上下文管理器之前，我们要清楚一个知识点：什么是生成式？

下面我直接引用官网的一段话（yield表达式）来描述生成器表达式：

yield 表达式在定义 generator 函数或是 asynchronous generator 的时候才会用到。 因此只能在函数定义的内部使用yield表达式。 在一个函数体内使用 yield 表达式会使这个函数变成一个生成器，并且在一个 async def 定义的函数体内使用 yield 表达式会让协程函数变成异步的生成器。

当一个生成器函数被调用的时候，它返回一个迭代器，称为生成器。然后这个生成器来控制生成器函数的执行。当这个生成器的某一个方法被调用的时候，生成器函数开始执行。这时会一直执行到第一个 yield 表达式，在此执行再次被挂起，给生成器的调用者返回 expression_list 的值。挂起后，我们说所有局部状态都被保留下来，包括局部变量的当前绑定，指令指针，内部求值栈和任何异常处理的状态。通过调用生成器的某一个方法，生成器函数继续执行。此时函数的运行就和 yield 表达式只是一个外部函数调用的情况完全一致。恢复后 yield 表达式的值取决于调用的哪个方法来恢复执行。 如果用的是 __next__() (通常通过语言内置的 for 或是 next() 来调用) 那么结果就是 None. 否则，如果用 send(), 那么结果就是传递给send方法的值。

所有这些使生成器函数与协程非常相似；它们 yield 多次，它们具有多个入口点，并且它们的执行可以被挂起。唯一的区别是生成器函数不能控制在它在 yield 后交给哪里继续执行；控制权总是转移到生成器的调用者。

在 try 结构中的任何位置都允许yield表达式。如果生成器在(因为引用计数到零或是因为被垃圾回收)销毁之前没有恢复执行，将调用生成器-迭代器的 close() 方法. close 方法允许任何挂起的 finally 子句执行。

当使用 yield from 时，它会将所提供的表达式视为一个子迭代器。 这个子迭代器产生的所有值都直接被传递给当前生成器方法的调用者。 通过 send() 传入的任何值以及通过 throw() 传入的任何异常如果有适当的方法则会被传给下层迭代器。 如果不是这种情况，那么 send() 将引发 AttributeError 或 TypeError，而 throw() 将立即引发所传入的异常。

当下层迭代器完成时，被引发的 StopIteration 实例的 value 属性会成为 yield 表达式的值。 它可以在引发 StopIteration 时被显式地设置，也可以在子迭代器是一个生成器时自动地设置（通过从子生成器返回一个值）。

我们看下迭代器有哪些方法：

请注意在生成器已经在执行时调用以下任何方法都会引发 ValueError 异常。

  • generator.__next__()

    开始一个生成器函数的执行或是从上次执行的 yield 表达式位置恢复执行。 当一个生成器函数通过 __next__() 方法恢复执行时，当前的 yield 表达式总是取值为 None。 随后会继续执行到下一个 yield 表达式，其 expression_list 的值会返回给 __next__() 的调用者。 如果生成器没有产生下一个值就退出，则将引发 StopIteration 异常。此方法通常是隐式地调用，例如通过 for 循环或是内置的 next() 函数。

  • generator.send(*value*)

    恢复执行并向生成器函数“发送”一个值。 value 参数将成为当前 yield 表达式的结果。 send() 方法会返回生成器所产生的下一个值，或者如果生成器没有产生下一个值就退出则会引发 StopIteration。 当调用 send() 来启动生成器时，它必须以 None 作为调用参数，因为这时没有可以接收值的 yield 表达式。

  • generator.throw(*type*[, *value*[, *traceback*]])

    在生成器暂停的位置引发 type 类型的异常，并返回该生成器函数所产生的下一个值。 如果生成器没有产生下一个值就退出，则将引发 StopIteration 异常。 如果生成器函数没有捕获传入的异常，或引发了另一个异常，则该异常会被传播给调用者。

  • generator.close()

    在生成器函数暂停的位置引发 GeneratorExit。 如果之后生成器函数正常退出、关闭或引发 GeneratorExit (由于未捕获该异常) 则关闭并返回其调用者。 如果生成器产生了一个值，关闭会引发 RuntimeError。 如果生成器引发任何其他异常，它会被传播给调用者。 如果生成器已经由于异常或正常退出则 close() 不会做任何事。

晦涩的理论知识我们暂时可以只需要了解到生成器可以在yield出暂停重入即可。

contextmanager 的源码剖析

有了上面生成器知识点的回顾，我们再看下contextmanager的源码

1
2
3
4
5
def contextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper

我们看到经过contextmanager修饰的函数最后被放到了生成器上下文管理对象中，并且将该对象返回。

为了简化我先把_GeneratorContextManager的结构写出来

1
2
3
4
5
6
7
8
9
10
11
12
13
class _GeneratorContextManager(_GeneratorContextManagerBase,
                               AbstractContextManager,
                               ContextDecorator):
    """Helper for @contextmanager decorator."""

    def _recreate_cm(self):

        return self.__class__(self.func, self.args, self.kwds)

    def __enter__(self):
      

    def __exit__(self, type, value, traceback):

其中的要打印一下_GeneratorContextManagerBase 在初始化的时候会初始化生成器函数

1
2
3
4
5
6
7
8
9
10
11
class _GeneratorContextManagerBase:
    """Shared functionality for @contextmanager and @asynccontextmanager."""

    def __init__(self, func, args, kwds):
        self.gen = func(*args, **kwds)
        self.func, self.args, self.kwds = func, args, kwds
        # Issue 19330: ensure context manager instances have good docstrings
        doc = getattr(func, "__doc__", None)
        if doc is None:
            doc = type(self).__doc__
        self.__doc__ = doc
1
2
3
4
5
6
7
8
9
10
11
@contextmanager
def some_generator(<arguments>):
    <setup>
    try:
        yield <value>
    finally:
        <cleanup>


with some_generator(<arguments>) as <variable>:
    <body>

当我们使用实现的上下文的时候首先会进入到__enter__协议函数

1
2
3
4
5
6
7
8
9
def __enter__(self):
    # do not keep args and kwds alive unnecessarily
    # they are only needed for recreation, which is not possible anymore
    del self.args, self.kwds, self.func
    try:
				# 执行生成器函数
        return next(self.gen)
    except StopIteration:
        raise RuntimeError("generator didn't yield") from None

我们知道next函数会直接执行生成器即yield之前的部分，当yield后面有值的时候会当做with as后面的变量。

当上面的过程咩有问题的时候就进入了with body，在执行with body代码的时候会有两种情况：有咩有异常发生。

我们先看正常退出with body之后将会执行__exit___的代码：

1
2
3
4
5
6
7
8
9
10
11
def __exit__(self, type, value, traceback):
    if type is None:
        try:
						
            # 再次执行生成器函数
            next(self.gen)

        except StopIteration:
            return False
        else:
            raise RuntimeError("generator didn't stop")

当再次执行生成器的时候将会执行yield后面的代码段。

不过因为我们对一个只有一个yield的生成器函数执行了两次yield。将会在执行完yield后面代码之后出发一个StopIteration异常，然后被捕捉到。

然后我们看下with body发生异常的时候

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def __exit__(self, type, value, traceback):
    if type is None:

    else:
        if value is None:
            # Need to force instantiation so we can reliably
            # tell if we get the same exception back
            value = type()
        try:
          	# 这里会在 生成器暂停的地方抛出 with body 里面发生的异常
            self.gen.throw(type, value, traceback)
        except StopIteration as exc:
            # Suppress StopIteration *unless* it's the same exception that
            # was passed to throw().  This prevents a StopIteration
            # raised inside the "with" statement from being suppressed.
            return exc is not value
        except RuntimeError as exc:
            # Don't re-raise the passed in exception. (issue27122)
            if exc is value:
                return False
            # Likewise, avoid suppressing if a StopIteration exception
            # was passed to throw() and later wrapped into a RuntimeError
            # (see PEP 479).
            if type is StopIteration and exc.__cause__ is value:
                return False
            raise
        except:
            # only re-raise if it's *not* the exception that was
            # passed to throw(), because __exit__() must not raise
            # an exception unless __exit__() itself failed.  But throw()
            # has to raise the exception to signal propagation, so this
            # fixes the impedance mismatch between the throw() protocol
            # and the __exit__() protocol.
            #
            # This cannot use 'except BaseException as exc' (as in the
            # async implementation) to maintain compatibility with
            # Python 2, where old-style class exceptions are not caught
            # by 'except BaseException'.
            if sys.exc_info()[1] is value:
                return False
            raise
        raise RuntimeError("generator didn't stop after throw()")

当with body发生异常之后将会执行到self.gen.throw(type, value, traceback)这句话。

我们知道这句话会在生成器暂停的位置引发 type 类型的异常，并返回该生成器函数所产生的下一个值。 如果生成器没有产生下一个值就退出，则将引发 StopIteration 异常。 如果生成器函数没有捕获传入的异常，或引发了另一个异常，则该异常会被传播给调用者。

上面就是通过contextmanager实现的上下文执行的完整流程。 我们看两个例子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
In [34]: @contextmanager
    ...: def catch_error():
    ...:
    ...:     try:
    ...:
    ...:         yield
    ...:
    ...:     except Exception as err:
    ...:
    ...:         print(err)
    ...:

In [35]: with catch_error():
    ...:     print("test")
    ...:
test

In [36]: with catch_error():
    ...:     print("test")
    ...:     1/0
    ...:
    ...:
test
division by zero

我们看到在with body中的异常抛到了生成器中并将其打印了出来。

这样我们使用contextmanager实现了一个具有捕获异常将其打印或者写入日志的上下文管理器，对于可能发生异常的代码比try except更加优美。

对于异步的asynccontextmanager的原理大致和同步的一致

参考文章：

python使用@contextmanager来定义上下文管理器

python黑魔法—上下文管理器

知识就是财富
如果您觉得文章对您有帮助, 欢迎请我喝杯水！