之前一直想写contextlib.contextmanager的实现原理，但也一直都没时间，今天看了Real Python关于上下文管理器的文章，打算补一下之前的坑。
预备知识
Python的语法糖很多，而contextlib.contextmanager主要涉及了三个比较常用的语法糖：装饰器、生成器和上下文管理器（with语句）。
装饰器
装饰器本质上原理很简单：它无非就是起到一种“锦上添花”的功能，在原有的函数上补充功能。被包装的函数以参数的形式传入到装饰器。也就是
@decorator
def func(): ...本质上等价于
func = decorator(func)生成器
生成器则可以看成能中途停下来（暂停的时候还能通过yield往外部抛出值），等待下一步指令并继续执行（外部可调用send方法传进指令，也可以直接通过next(...)让它继续）的函数。这里不再赘述，直接来个简单的例子：
def gen():
print("gen(): Start...")
first_command = yield "First"
print(f"gen(): {first_command = }")
second_command = yield "Second"
print(f"gen(): {second_command = }")
if __name__ == '__main__':
g = gen()
print(f"recv1: {next(g)}")
print(f"recv2: {g.send('Command1')}")
try:
print(f"recv3: {g.send('Command2')}") # This will raise StopIteration
except StopIteration:
pass可以看到，我们首先调用生成器函数gen得到一个生成器迭代器g，然后使用next(g)激活生成器迭代器g（此时会输出"gen(): Start..."），激活后，会返回一个字符串"First"。然后，我们使用g.send('Command1')发送字符串'Command1'到生成器迭代器中，此时生成器迭代器内的first_command = yield "First"收到外部传入的参数并赋值first_command，也就是first_command='Command1'，并继续执行下去。当生成器迭代器到达终点时，其会抛出一个StopIteration异常告知外部已经。
需要注意的是，如果外部不使用
send发送参数而是直接调用next，则生成器迭代器内通过yield表达式所收到的值为None。
这里厘清一下生成器相关的术语：
- 生成器函数（generator function，大多数场合下，生成器（generator）指的是生成器函数）：函数体内包含
yield语句的函数，调用它会返回一个生成器迭代器。 - 生成器迭代器（generator iterator）：生成器函数所创建的对象，其实现了迭代器协议（
__iter__和__next__方法），具有send、throw、close等方法。
一般来说，生成器指的是生成器函数，但有时候也会指生成器迭代器，具体要看上下文。
上下文管理器
上下文管理器则是实现了__enter__和__exit__这两个魔术方法的类，使用时，将该类的实例传入到with语句即可。
我们可以简单实现一个上下文管理器，其会在进出上下文的时候输出一些信息，并且捕获with语句体内的ZeroDivisionError异常。
class SimpleCtxManager:
def __enter__(self):
print("Entering the context")
return self
def __exit__(self, exc_type, exc_value, traceback):
print("Exiting the context")
print(f"{exc_type = }")
print(f"{exc_value = }")
print(f"{traceback = }")
if isinstance(exc_value, ZeroDivisionError):
print("Handled ZeroDivisionError")
return True # Suppress the exception
return False # Propagate other exceptions我们可以写个简单的使用例子：
if __name__ == '__main__':
with SimpleCtxManager() as s:
print(f"{s = }")
print("OK, no exception\n")
with SimpleCtxManager() as cm:
1 / 0 # This will raise ZeroDivisionError
print("Go ahead, no exception propagated\n")
with SimpleCtxManager() as cm:
a = []
a[1] = -1
print("This line will not be printed due to IndexError")运行程序，可以看到无论是否有异常，上下文管理器都能执行到__exit__方法，从而确保其所托管的资源能释放。当然，如果异常继续传播，with语句之后的代码也不会执行了（除非遇到finally语句）。
Entering the context
s = 'Hello'
Exiting the context
exc_type = None
exc_value = None
traceback = None
OK, no exception
Entering the context
Exiting the context
exc_type = <class 'ZeroDivisionError'>
exc_value = ZeroDivisionError('division by zero')
traceback = <traceback object at 0x000001AA8C7CFA40>
Handled ZeroDivisionError
Go ahead, no exception propagated
Entering the context
Exiting the context
exc_type = <class 'IndexError'>
exc_value = IndexError('list assignment index out of range')
traceback = <traceback object at 0x000001AA8C7CF9C0>
python-BaseException
Traceback (most recent call last):
File "C:\Users\Jeza\PyQtEx\test1.py", line 43, in <module>
a[1] = -1
~^^^
IndexError: list assignment index out of rangecontextmanager的简单用法
主角终于登场，现在我们使用contextlib.contextmanager对上文的上下文管理器类SimpleCtxManager做个等价的改造：
@contextlib.contextmanager
def SimpleCtxManager():
print("Entering the context")
try:
yield
except Exception as e:
print(f"exc_type = {type(e)}")
print(f"exc_value = {e}")
print(f"traceback = {e.__traceback__}")
if isinstance(e, ZeroDivisionError):
print("Handled ZeroDivisionError:", e)
else:
raise
finally:
print("Exiting the context")可以看到，contextlib.contextmanager是一个装饰器，其所包装的是一个生成器函数，其中yield语句前面的代码可以看成__enter__的逻辑，而后面的except…finally…语句则可以看成__exit__的逻辑（处理异常/清理资源）。对比发现，使用contextlib.contextmanager 不仅能让代码易读（看起来更流程化），而且不用专门写一个类，仅需一个生成器函数即可做到一样的功能。
contextmanager的源码分析
OK，终于讲到contextlib.contextmanager的原理了。我们可以在Python官方标准库中找到contextlib.py的源代码，定位到contextmanager的实现：
def contextmanager(func):
@wraps(func)
def helper(*args, **kwds):
return _GeneratorContextManager(func, args, kwds)
return helper由代码看出，contextmanager是一个装饰器，其接受一个函数func作为参数，在闭包helper内用它来继续构造上下文管理器类_GeneratorContextManager的实例并返回。我们将装饰器contextmanager所返回的闭包helper视作为一个上下文管理器工厂（在外部看起来，它其实用法和上下文管理器是一样的），其接受任意的参数，负责在使用处构造出真正的上下文管理器_GeneratorContextManager的实例。
三个基类的分析
我们继续研究_GeneratorContextManager类的实现，首先看看它的基类：
class _GeneratorContextManager(
_GeneratorContextManagerBase,
AbstractContextManager,
ContextDecorator,
):
...可以看到其继承了三个基类，我们首先看第一个基类_GeneratorContextManagerBase的实现，这也是所有基于生成器的上下文管理器的公共基类：
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
# Unfortunately, this still doesn't provide good help output when
# inspecting the created context manager instances, since pydoc
# currently bypasses the instance docstring and shows the docstring
# for the class instead.
# See http://bugs.python.org/issue19404 for more details.
def _recreate_cm(self):
# _GCMB instances are one-shot context managers, so the
# CM must be recreated each time a decorated function is
# called
return self.__class__(self.func, self.args, self.kwds)不难分析，这个基类所实现的逻辑比较简单：保存生成器函数func以及调用func所需的参数args，kwds 、将func的__doc__属性复制到生成器里面去，以及提供一个_recreate_cm方法，以用来重新创建一个新的上下文管理器实例，保证该实例作为装饰器使用时，在每一次函数调用都能秽土重生，重新生成一个全新的实例用来执行上下文——因为生成器迭代器是一次性的，一旦激活并next下去，就没有回头路了。需要回头？只能重来。（让生成器函数再生一个出来）
再次强调下，这里的
_recreate_cm方法仅用于该上下文管理器实例作为装饰器使用的场合。其覆写了第三个基类ContextDecorator的同名方法，后面会讲到。
这个上下文管理器的实例如何作为装饰器使用呢？后面第三个基类
ContextDecorator的分析会有例子。这两个基类的关系很密切。
OK，我们来看第二个基类AbstractContextManager的实现：
class AbstractContextManager(abc.ABC):
"""An abstract base class for context managers."""
__class_getitem__ = classmethod(GenericAlias)
def __enter__(self):
"""Return `self` upon entering the runtime context."""
return self
@abc.abstractmethod
def __exit__(self, exc_type, exc_value, traceback):
"""Raise any exception triggered within the runtime context."""
return None
@classmethod
def __subclasshook__(cls, C):
if cls is AbstractContextManager:
return _collections_abc._check_methods(C, "__enter__", "__exit__")
return NotImplemented这个没什么好说的，其就是所有上下文管理器的基类。当然，由于Python是鸭子类型语言，并不强制上下文管理器的实现必须继承这个基类，这个基类只是用来加强类型检查，避免__enter__或__exit__方法没有实现（即上述的__subclasshook__所做的）。比如在下面的例子中，我们定义WrongCtxManager类的时候，犯了个很致命的错误——只实现了__enter__却漏了__exit__——还好！由于这个类继承了AbstractContextManager，在实例化的时候Python解释器会“贴心”地抛出异常，告知你没有实现完整的上下文管理协议。
from contextlib import AbstractContextManager
class WrongCtxManager(AbstractContextManager):
def __enter__(self):
print("Entering the context")
return self
cm = WrongCtxManager() # error!可见当鹅比当鸭好！
我们来看第三个基类ContextDecorator：
class ContextDecorator(object):
"A base class or mixin that enables context managers to work as decorators."
def _recreate_cm(self):
"""Return a recreated instance of self.
Allows an otherwise one-shot context manager like
_GeneratorContextManager to support use as
a decorator via implicit recreation.
This is a private interface just for _GeneratorContextManager.
See issue #11647 for details.
"""
return self
def __call__(self, func):
@wraps(func)
def inner(*args, **kwds):
with self._recreate_cm():
return func(*args, **kwds)
return inner这个基类更像是一个mixin，为子类实现了一个__call__魔术方法，使得上下文管理器实例可作为装饰器使用。在被包装函数调用的前后，执行__enter__和__exit__方法。关于_recreate_cm，前面已经提前路透了，这里不再多说了。
以前面的SimpleCtxManager为例，我们为它加上父类ContextDecorator，此时其实例就可以作为装饰器来用了。
from contextlib import ContextDecorator
class SimpleCtxManager(ContextDecorator): # Inherit from ContextDecorator
def __enter__(self):
print("Entering the context")
return self
def __exit__(self, exc_type, exc_value, traceback):
print("Exiting the context")
print(f"{exc_type = }")
print(f"{exc_value = }")
print(f"{traceback = }")
if isinstance(exc_value, ZeroDivisionError):
print("Handled ZeroDivisionError")
return True # Suppress the exception
return False # Propagate other exceptions
cm = SimpleCtxManager()
@cm # Use the instance as a decorator
def func():
print("func...")
func()
# Outputs:
# Entering the context
# func...
# Exiting the context
# exc_type = None
# exc_value = None
# traceback = None回到contextlib.contextmanager
好了，分析完三个基类后，我们终于轮到分析主角contextlib.contextmanager了，可以看到，通过继承上述的三个基类后，contextlib.contextmanager只需将精力花在最核心的__enter__和__exit__两个方法：
（为减少篇幅，下面的源码去掉了注释）
class _GeneratorContextManager(
_GeneratorContextManagerBase,
AbstractContextManager,
ContextDecorator,
):
def __enter__(self):
del self.args, self.kwds, self.func
try:
return next(self.gen)
except StopIteration:
raise RuntimeError("generator didn't yield") from None
def __exit__(self, typ, value, traceback):
if typ is None:
try:
next(self.gen)
except StopIteration:
return False
else:
try:
raise RuntimeError("generator didn't stop")
finally:
self.gen.close()
else:
if value is None:
value = typ()
try:
self.gen.throw(typ, value, traceback)
except StopIteration as exc:
return exc is not value
except RuntimeError as exc:
if exc is value:
exc.__traceback__ = traceback
return False
if (
isinstance(value, StopIteration)
and exc.__cause__ is value
):
value.__traceback__ = traceback
return False
raise
except BaseException as exc:
if exc is not value:
raise
exc.__traceback__ = traceback
return False
try:
raise RuntimeError("generator didn't stop after throw()")
finally:
self.gen.close()可以看到，_GeneratorContextManager是幕后真正起作用的上下文管理器。我们抛开一些细节性的校验逻辑，分析下__enter__的主要思路：使用next(self.gen)激活生成器迭代器（也就是contextmanager包裹的生成器函数所返回的生成器迭代器），使其走到yield语句处停下——这也意味着，yield语句前面所执行的代码拿来充当__enter__的逻辑了！
我们再看__exit__的实现思路，分两种情况：
- 如果
with语句块没有抛出异常：使用next(self.gen)恢复生成器迭代器，并期望它能执行到最后（包括走完finally语句块）。如果发现生成器迭代器暂停了（说明又遇到了yield），则这个生成器函数的实现不符合contextmanager的规范，抛出RuntimeError。 - 如果
with语句块抛出了异常：将异常转发到生成器迭代器的yield表达式处（self.gen.throw(typ, value, traceback)），给它一个机会做清理或选择抑制异常，然后再根据它所抛出的异常再做处理——- 如果抛出的是
StopIteration且其异常实例不同于所传入的异常实例（确保了这个StopIteration是生成器迭代器自己执行完毕所抛出的，而非重新raise了所传入的异常（尤其这个所传的异常也是StopIteration）），意味着它选择了抑制异常，此时__exit__返回True； - 如果抛出的是其他异常，则说明生成器选择了不抑制异常，此时
__exit__返回False，让异常继续传播。
- 如果抛出的是
里面还有一些细节也挺有意思的，这里就先不讲了，有空的话再原地补坑。先这样，睡觉~
