上一篇文章中，我們介紹了 python 中的迭代器與生成器。 python 中的迭代器與生成器
此前的文章中，我們已經看過上下文管理器的例子。 python 魔術方法（四）非常用方法與運算符重載方法大合集
本文我們通過分析標准庫中 contextlib.contextmanager 裝飾器的源碼，來看看如何讓他們結合起來生成更加優雅的代碼。
class Test:
def __enter__(self):
print('now in __enter__')
return 'Hello World'
def __exit__(self, exc_type, exc_val, exc_tb):
print('now exit')
return True
if __name__ == '__main__':
test = Test()
with test as teststr:
print(teststr)
print('end of main')調用打印出了：
now in __enter__ Hello World now exit end of main
當 with 塊被執行時，解釋器會自動調用對象的 __enter__ 方法。 而在 with 塊結束時，解釋器則會自動調用對象的 __exit__ 方法，__exit__ 方法最終可以選擇返回 True 或拋出異常。
標准庫中，contextlib.contextmanager 裝飾器通過 yield 關鍵字可以減少創建上下文管理器的樣板代碼量。 上面的例子可以改造為：
import contextlib
class Test:
@contextlib.contextmanager
def contextmanager(self):
print('now in __enter__')
yield 'Hello World'
print('now exit')
return True
if __name__ == '__main__':
test = Test()
with test.contextmanager() as teststr:
print(teststr)
print('end of main')同樣打印出了：
now in __enter__ Hello World now exit end of main
本質上 contextlib.contextmanager 仍然是利用了 yield 生成器的特性，他將函數包裝並增加了 __enter__ 與 __exit__ 兩個方法。
def contextmanager(func):
@wraps(func)
def helper(*args, **kwds):
return _GeneratorContextManager(func, args, kwds)
return helper
class _GeneratorContextManager():
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
def __enter__(self):
try:
return next(self.gen)
except StopIteration:
raise RuntimeError("generator didn't yield") from None
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")
else:
if value is None:
# Need to force instantiation so we can reliably
# tell if we get the same exception back
value = type()
try:
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.
#
if sys.exc_info()[1] is value:
return False
raise
raise RuntimeError("generator didn't stop after throw()")可以看到，__enter__ 方法實現的比較簡單，僅僅是通過 next 方法獲取了生成器的首個生成的數據。 __exit__ 方法則相對復雜：
從上述代碼我們可以看到一個嚴重的問題：__enter__ 代碼是未捕獲異常的，一旦我們在 with 塊中拋出異常，則會導致 __exit__ 中的清理代碼無法被執行。
import contextlib
class Test:
@contextlib.contextmanager
def contextmanager(self):
print('now in __enter__')
yield self.raiseexc(1)
print('now exit')
return True
def raiseexc(self, param):
if param < 5:
raise Exception('test exception')
if __name__ == '__main__':
test = Test()
with test.contextmanager() as teststr:
print(teststr)
print('end of main')執行，打印出了：
now in __enter__ Traceback (most recent call last): File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydevd.py", line 1741, in <module> main() File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydevd.py", line 1735, in main globals = debugger.run(setup[‘file’], None, None, is_module) File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydevd.py", line 1135, in run pydev_imports.execfile(file, globals, locals) # execute the script File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydev_imps\_pydev_execfile.py", line 18, in execfileexec(compile(contents+"\n", file, ’exec’), glob, loc)File "D:/Workspace/code/python/testpython/fluentpython/contextmanager.py", line 19, in <module>with test.contextmanager() as teststr:File "C:\ProgramData\Anaconda3\lib\contextlib.py", line 81, in \_enter__ return next(self.gen) File "D:/Workspace/code/python/testpython/fluentpython/contextmanager.py", line 8, in contextmanager yield self.raiseexc(1) File "D:/Workspace/code/python/testpython/fluentpython/contextmanager.py", line 14, in raiseexc raise Exception(‘test exception’) Exception: test exception
所以，在使用 @contextlib.contextmanager 時千萬要注意，不能在 yield 執行時拋出異常。
import contextlib
class Test:
@contextlib.contextmanager
def contextmanager(self):
print('now in __enter__')
try:
yield self.raiseexc(1)
except Exception:
print('exception happened')
print('now exit')
return True
def raiseexc(self, param):
if param < 5:
raise Exception('test exception')
if __name__ == '__main__':
test = Test()
with test.contextmanager() as teststr:
print(teststr)
print('end of main')打印出了：
now in __enter__ exception happened now exit Traceback (most recent call last): File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydevd.py", line 1741, in <module> main() File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydevd.py", line 1735, in main globals = debugger.run(setup[‘file’], None, None, is_module) File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydevd.py", line 1135, in run pydev_imports.execfile(file, globals, locals) # execute the script File "C:\Program Files\JetBrains\PyCharm 2018.3.1\helpers\pydev\pydev_imps\_pydev_execfile.py", line 18, in execfileexec(compile(contents+"\n", file, ’exec’), glob, loc)File "D:/Workspace/code/python/testpython/fluentpython/contextmanager.py", line 22, in <module>with test.contextmanager() as teststr:File "C:\ProgramData\Anaconda3\lib\contextlib.py", line 83, in \_enter__ raise RuntimeError("generator didn’t yield") from None RuntimeError: generator didn’t yield
雖然仍然拋出了異常，但我們看到 __exit__ 方法中的清理代碼仍然得以被執行。
