今天我们来看看contextlib这个库里面的contextmanagerclosing怎么用,为什么这么用,背后的原理是什么。首先看看contextmanager

contextmanager

使用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from contextlib import contextmanager
class Query(object):
def __init__(self, name):
self.name = name
def query(self):
print('Query info about %s...' % self.name)
@contextmanager
def create_query(name):
print('Begin')
q = Query(name)
yield q
print('End')

这样就定义了一个上下文管理器,create_query也可以使用with语句了,比如这样:

1
2
with create_query('Bob') as q:
q.query()

with语句首先执行yield之前的语句,yield调用会返回后面的变量,然后执行with语句内部的所有语句,with内部语句执行完成之后执行yield之后的语句

所以我们就可以在yield之前定义类似__enter__的代码,yield后面的变量就是__enter__的返回值,然后yield之后的语句是定义类似__exit__的代码。

清楚了使用,来看看源代码:

源代码

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
"""@contextmanager decorator.
Typical usage:
@contextmanager
def some_generator(<arguments>):
<setup> # 初始化代码 __enter__
try:
yield <value> # __enter__返回的值
finally:
<cleanup> # __exit__
This makes this:
with some_generator(<arguments>) as <variable>:
<body>
equivalent to this:
<setup>
try:
<variable> = <value>
<body>
finally:
<cleanup>
"""

其代码是调用了一个_GeneratorContextManager

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

相当于用contextmanager装饰的函数在调用的时候,比如with func() as f:的时候,func()其实是调用了helper(),而这个返回的是一个_GeneratorContextManager对象,也就是说其实是_GeneratorContextManager对象实现了上下文协议,这个对象肯定就定义了__enter____exit__方法,当然这个方法是与func有关的,那么继续来看看_GeneratorContextManager

_GeneratorContextManager

这里有一个要点就是使用类来实现了一个装饰器,这种做法被推荐过,因为可以更好的弄清楚内部变量的作用域之类的东西。

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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
class _GeneratorContextManager(ContextDecorator):
"""Helper for @contextmanager decorator. 发现helper这个经常用,啊哈,以后我也开始用这个变量名"""
def __init__(self, func, args, kwds):
# func是一个生成器函数,这里首先"实例化"
self.gen = func(*args, **kwds)
# 保存之前函数的一些参数
self.func, self.args, self.kwds = func, args, kwds
doc = getattr(func, "__doc__", None)
if doc is None:
doc = type(self).__doc__
self.__doc__ = doc
def _recreate_cm(self):
# _GCM 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)
def __enter__(self):
try:
# enter方法,返回的肯定是生成器yield的值
# 在next执行的过程中,已经将yield之前的代码执行了
return next(self.gen)
except StopIteration:
raise RuntimeError("generator didn't yield") from None
def __exit__(self, type, value, traceback):
# 如果没有错误,继续next生成器
# 如果生成器没有结束,那么引发异常
if type is None:
try:
next(self.gen)
except StopIteration:
return
else:
raise RuntimeError("generator didn't stop")
# 如果在with语句中有异常
else:
if value is None:
# 实例化异常类型的实例
value = type()
try:
# 向生成器发送这个异常,使其停止
self.gen.throw(type, value, traceback)
# 如果没有停止,则引发RuntimeError异常
raise RuntimeError("generator didn't stop after throw()")
except StopIteration as exc:
# 生成器停止异常,如果exit处理的异常是StopIteration,则返回True,相当于已经执行了清理代码
# 如果exit处理的异常不是停止异常,则继续向上引发这个异常,说明这个异常已经使得清理代码执行,但是需要往上引发
return exc is not value
except RuntimeError as exc:
# 如果在将exit的异常(比如是生成器停止异常)发送给生成器,但是没有停止,所以捕捉到了RuntimeError异常
# 如果该RuntimeError的异常是由于exit处理的异常(生成器停止异常)引发的,则继续引发这个异常,因为是先发生生成器停止异常,然后再发生RuntimeError,所以会继续向上引发RuntimeError
if exc.__cause__ is value:
return False
# 否则向上引发RuntimeError
raise
except:
# 只会重新引发不是传递给throw的异常,因为__exit__一定不会引发异常除非__exit__本身失败了
# 但是throw方法不得不引发异常导致异常传播,因此这里修正了throw协议和__exit__协议的不匹配问题
# 也就是说,这个异常一定不能是throw发出的,因为__exit__不会引发异常的,所以只有当这个处理的异常不是throw发出的,即exit传递进来的value,就会重新引发
if sys.exc_info()[1] is not value:
raise

有一点需要说明的是__exit__方法返回True,说明已经处理好了,返回False,会继续抛出异常。

这里有一大堆的错误处理,以后可以参考。要理解其核心点在于,exit方法接收到的异常可能是StopIteration,因为可能在with语句中将生成器迭代完成了,比如这样:

1
2
with create_query('Bob') as q:
next(q)

当然形式可能不一样,但是效果都是在with中引发了StopIteration异常,然后这里处理的时候就会有这么多种情况,异常处理的主要目的是抑制StopIteration异常的发出。

closing

还有一个简单的上下文处理器就是这个closing,主要的功能就是给加了一个__exit__方法和__enter__方法,使其可以支持上下文管理器协议,下面是使用方法:

1
2
3
4
5
6
from contextlib import closing
from urllib.request import urlopen
with closing(urlopen('https://www.python.org')) as page:
for line in page:
print(line)

下面是源码:

1
2
3
4
5
6
7
class closing(object):
def __init__(self, thing):
self.thing = thing
def __enter__(self):
return self.thing
def __exit__(self, *exc_info):
self.thing.close()

也可以这么实现,很机智吧!

1
2
3
4
5
6
@contextmanager
def closing(thing):
try:
yield thing
finally:
thing.close()

小结

contextlib库常用的内容就是这些了,不得不感叹设计真的是巧妙,非常不错。