在 Part 1 我们已经介绍了生成器的定义和生成器的操作,现在让我们开始使用生成器。Part 2 主要描述了如何使用 yield
和 contextmanager
创建一个上下文管理器,并解释了原理。
上下文管理器
理解上下文可以联想我们做阅读理解时要解读文章某处的意思需要阅读该处前后段落,正是前后文提供了理解的“背景”。而程序的运行的上下文也可以理解为程序在运行时的某些变量,正是这些变量构成了运行环境,让程序可以完成操作。
Python 中的上下文管理器提供这样一种功能,为你的程序运行时提供一个特定“空间”,当进入这个空间时 Python 上下文管理器 为你做一些准备工作。这个“空间”中一般含有特殊的变量,你在这个“空间”中进行一些操作,然后离开。在你离开时 Python 上下文管理器又会帮你做一些收尾工作,保证不会污染运行环境。
下面是一些常见的代码模式
# 读取文件
f = open()
# do something
f.close()
# 使用锁
lock.acquire()
# do somethin
lock.release()
# 进行数据库操作
db.start_transaction()
# do something
db.commit()
# 对某段代码进行计时
start = time.time()
# do something
end = time.time()
这些代码进行的都是“先做这个(准备工作,比如获取一个数据库连接),然后做这个(比如写入数据),最后整理工作环境(如提交改动,关闭链接,释放资源等)。
如果使用 with
可以这样写:
witn open(filename) as f:
# do something
pass
with lock():
# do something
pass
with
语句实际上使用了实现了 __enter__
和 __exit__
方法的上下文管理器类。一个典型的上下文管理器类如下:
clss ContextManager:
def __enter__(self):
return value
def __exit__(self, exc_type, val, tb):
if exec_type is None:
return
else:
# 处理异常
return True if handled else False
正如方法名明确告诉我们的,__enter__
方法负责进入上下的准备工作,如果有需要可以返回一个值,这个值将会被赋值给 with ContextManager() as ret_value
中的 ret_value
。__exit__
则负责收尾工作,这包括了异常处理。
对于这样一段代码
with ContextManager() as var:
# do something
相当于
ctxmanager = ContextManager()
var = ctxmanager.__enter__()
# do somethin
ctxmanager.__exit__()
一个可用的例子:
import tempfile
import shutil
class TmpDir:
def __enter__(self):
self.dirname = tempfile.mkdtemp()
return self.dirname
def __exit__(self, exc, val, tb):
shutil.rmtree(self.dirname)
这个上下文管理提供临时文件的功能,在 with
语句结束后会自动删除临时文件夹。
with TempDir() as dirname:
# 使用临时文件夹进行一些操作
pass
关于上面两个特殊方法的文档可以在 Python 文档的 Context Manager Types 找到。另外关于 with
关键字的详细说明参考 PEP 343,不过这篇 PEP 不是很好读,Good Luck :simple_smile:!
使用 yield 和 contextmanager
能看到这里的都应该对上下文管理器有所了解,准备好把 yield
加入我们的上下文管理器代码中。
先看一个例子
import tempfile, shutil
from contextlib import contextmanager
@contextmanager
def tempdir():
outdir = tempfile.mkdtemp()
try:
yield outdir
finally:
shutil.rmtree(outdir)
与使用上下文管理器类的实现方式不同,这里我们没有显式实现 __enter__
和 __exit__
,而是通过 contextmanager
装饰器和 yield
实现,你可以试试这两种方式是等价的。
要理解上面的代码,可以把 yield
想象为一把剪刀,把这个函数一分为二,上部分相当于 __enter__
,下部分相当于 __exit__
。我这样说大家应该明白了吧。
import tempfile, shutil
from contextlib import contextmanager
@contextmanager
def tempdir():
outdir = tempfile.mkdtemp() #
try: # __enter__
yield outdir #
--cut---╳-----------------------------------
finally: #
shutil.rmtree(outdir) # __exit__
实现“剪刀”功能关键在于 contextmanager
。对于上面的代码,我们来一步一步地结构它:
contextmanager 装饰器
contextmanager
其实使用了一个上下文管理器类,这个类在在初始化时需要提供一个生成器。
class GeneratorCM:
def __init__(self, gen):
self.gen = gen
def __enter__(self):
...
def __exit__(self, exc, val, tb):
...
contextmanager
的实现如下
def contextmanager(func):
def run(*args, **kwargs):
return GeneratorCM(func(*args, **kwargs))
return run
由于 contextmanger
所装饰的函数里有 yield
所以我们在调用 func(*args, **kwargs)
时返回的是一个生成器。要使这个生成器前进,我们需要调用 next
函数
让生成器前进
def __enter__(self):
return next(self.gen)
GeneratorCM
的 __ente__
方法会让生成器前进到 yield
语句处,并返回产出值。
收尾
def __exit__(self, exc, val, tb):
try:
if exc is None:
next(self.gen)
else:
self.gen.throw(exc, val, tb)
raise RuntimeError('Generator didn\'t stop')
except StopIteration:
return True
except:
if sys.exc_info()[1] is not val: raise
__exit__
函数的逻辑比较复杂,如果没有传入异常,首先它会尝试对生成器调用 next
,正常情况下这会抛出 StopIteration
,这个异常会被不做并返回 True
,告诉解释器正常退出;如果传入异常,会使用 throw
在 yield
处抛出这个异常;如果有其他未捕捉的错误,就重新抛出该错误。
实际的代码实现会更加复杂,还有一些异常情况没有处理
- 没有相关值的异常
- 在 with 语句块中抛出的
StopIteration
- 在上下文管理器中抛出的异常
如果你对怎么实现感兴趣,你可以阅读代码或者再一次阅读 PEP 343。
总结
Part 2 都是关于上下文管理器的内容,与协程关系不大。但通过这部分我们可以看到 yield
完全不同的用法,也熟悉了控制流 (control-flow) ,这与 Part 3 的异步处理流程有很大关系。让我们 Part 3 再见。