(轉)contextlib — 上下文管理器工具

本節目標: 創建和使用基於上下文管理器的工具

contextlib 模塊包含了可使用 with 語句的上下文管理工具，一起了解一下。

上下文管理器 API##

上下文管理器 就是一個給包含在其中代碼塊提供資源的對象，在進入塊時創建一些資源，在退出塊后清理掉。舉個例子，文件操作就支持上下文管理器 API，使用這種方法就能保證在讀完寫完后總能關閉文件，並且寫起來很簡單。

contextlib_file.py

with open('/tmp/pymotw.txt', 'wt') as f: f.write('contents go here') # 運行到這文件就自動關閉了。

每個上下文管理器都允許使用 with 語句來執行，在寫的時候也要包含兩個必須的方法。 __enter__() 方法是 with 進入代碼時所執行的方法，一般要返回一個對象以讓處在代碼塊中的代碼使用。離開 with 代碼塊后，上下文管理器中的 __exit__() 方法就會被調用以清理一些用過的資源。

contextlib_api.py

class Context: def __init__(self): print('__init__()') def __enter__(self): print('__enter__()') return self def __exit__(self, exc_type, exc_val, exc_tb): print('__exit__()') with Context(): print('Doing work in the context')

結合 with 與上下文管理器可以書寫非常緊湊舒服的 try:finally 代碼塊，因為上下文管理器中的 __exit__() 無論如何都一定會被執行，即使有異常拋出。

$ python3 contextlib_api.py __init__() __enter__() Doing work in the context __exit__()

__enter__() 方法可以返回任意對象，它返回的任何對象都會被賦給 with 語句中的 as 所指向的變量。本例中可以看到 Context 返回了一個在之后使用的對象。

contextlib_api_other_object.py

class WithinContext: def __init__(self, context): print('WithinContext.__init__({})'.format(context)) def do_something(self): print('WithinContext.do_something()') def __del__(self): print('WithinContext.__del__') class Context: def __init__(self): print('Context.__init__()') def __enter__(self): print('Context.__enter__()') return WithinContext(self) def __exit__(self, exc_type, exc_val, exc_tb): print('Context.__exit__()') with Context() as c: c.do_something()

變量 c 就是 __enter__() 所返回的值，並不一定是 with 語句創建的 Context 實例才可以使用上下文管理器，外部創建的實例同樣可以使用 with 。

$ python3 contextlib_api_other_object.py Context.__init__() Context.__enter__() WithinContext.__init__(<__main__.Context object at 0x101e9c080>) WithinContext.do_something() Context.__exit__() WithinContext.__del__

__exit__() 方法所接受的參數是任何在 with 代碼塊中產生的異常的詳細信息。

contextlib_api_error.py

class Context: def __init__(self, handle_error): print('__init__({})'.format(handle_error)) self.handle_error = handle_error def __enter__(self): print('__enter__()') return self def __exit__(self, exc_type, exc_val, exc_tb): print('__exit__()') print(' exc_type =', exc_type) print(' exc_val =', exc_val) print(' exc_tb =', exc_tb) return self.handle_error with Context(True): raise RuntimeError('error message handled') print() with Context(False): raise RuntimeError('error message propagated')

如果上下文管理器可以處理這個異常， __exit__() 應該返回 True 表示這個異常並沒有造成麻煩，不必管它。如果返回的是 False，則該異常會在 __exit__() 執行后重新拋出。

$ python3 contextlib_api_error.py __init__(True) __enter__() __exit__() exc_type = <class 'RuntimeError'> exc_val = error message handled exc_tb = <traceback object at 0x1044ea648> __init__(False) __enter__() __exit__() exc_type = <class 'RuntimeError'> exc_val = error message propagated exc_tb = <traceback object at 0x1044ea648> Traceback (most recent call last): File "contextlib_api_error.py", line 34, in <module> raise RuntimeError('error message propagated') RuntimeError: error message propagated

Cyrbuzz 翻譯於

由 Summer 審閱

函數裝飾器方式的上下文管理器##

ContextDecorator 類可以讓標准的上下文管理器類變成一個可以作為函數裝飾器方式使用的上下文管理器。

contextlib_decorator.py

import contextlib class Context(contextlib.ContextDecorator): def __init__(self, how_used): self.how_used = how_used print('__init__({})'.format(how_used)) def __enter__(self): print('__enter__({})'.format(self.how_used)) return self def __exit__(self, exc_type, exc_val, exc_tb): print('__exit__({})'.format(self.how_used)) @Context('as decorator') def func(message): print(message) print() with Context('as context manager'): print('Doing work in the context') print() func('Doing work in the wrapped function')

把上下文管理器作為函數裝飾器使用的不同之處在於 __enter__() 所返回的值無法在被裝飾的函數中使用，不能像 with 和 as一樣。被裝飾后的函數的參數仍像之前一樣傳遞。

$ python3 contextlib_decorator.py __init__(as decorator) __init__(as context manager) __enter__(as context manager) Doing work in the context __exit__(as context manager) __enter__(as decorator) Doing work in the wrapped function __exit__(as decorator)

從生成器到上下文管理器##

創建一個上下文管理器的傳統方式是寫一個類，然后寫它的 __enter__() 和 __exit__() 方法，這並不難寫，不過有時候把所有的東西都寫全並沒有必要。在這種情況下，可以使用 contextmanager() 裝飾器將一個生成器函數變成一個上下文管理器。

contextlib_contextmanager.py

import contextlib @contextlib.contextmanager def make_context(): print(' entering') try: yield {} except RuntimeError as err: print(' ERROR:', err) finally: print(' exiting') print('Normal:') with make_context() as value: print(' inside with statement:', value) print('\nHandled error:') with make_context() as value: raise RuntimeError('showing example of handling an error') print('\nUnhandled error:') with make_context() as value: raise ValueError('this exception is not handled')

生成器應首先初始化上下文，確保只生成一次，最后應清理上下文內容。所生成的內容可以被 with 語句的 as 所賦值給一個變量。 with 語句中的異常也會在生成器內部重新被拋出，這樣在我們就可以處理這個異常。

$ python3 contextlib_contextmanager.py Normal: entering inside with statement: {} exiting Handled error: entering ERROR: showing example of handling an error exiting Unhandled error: entering exiting Traceback (most recent call last): File "contextlib_contextmanager.py", line 33, in <module> raise ValueError('this exception is not handled') ValueError: this exception is not handled

contextmanager() 所返回的上下文管理器繼承自 ContextDecorator，所以同樣可以作為函數裝飾器來使用。

contextlib_contextmanager_decorator.py

import contextlib @contextlib.contextmanager def make_context(): print(' entering') try: # 通過 Yield 控制，但無需返回值，因為作為裝飾器 # 使用時，上下文管理器所返回的值並不會被使用到。 yield except RuntimeError as err: print(' ERROR:', err) finally: print(' exiting') @make_context() def normal(): print(' inside with statement') @make_context() def throw_error(err): raise err print('Normal:') normal() print('\nHandled error:') throw_error(RuntimeError('showing example of handling an error')) print('\nUnhandled error:') throw_error(ValueError('this exception is not handled'))

與上面 ContextDecorator 的例子一樣，上下文管理器作為裝飾器使用時生成的值並不能被被裝飾的函數所用。當然這種方式下本來的參數還是可以正常傳遞的，上面是以 throw_error() 為例演示的。

$ python3 contextlib_contextmanager_decorator.py Normal: entering inside with statement exiting Handled error: entering ERROR: showing example of handling an error exiting Unhandled error: entering exiting Traceback (most recent call last): File "contextlib_contextmanager_decorator.py", line 43, in <module> throw_error(ValueError('this exception is not handled')) File ".../lib/python3.6/contextlib.py", line 52, in inner return func(*args, **kwds) File "contextlib_contextmanager_decorator.py", line 33, in throw_error raise err ValueError: this exception is not handled

Cyrbuzz 翻譯於

由 Summer 審閱

關閉打開的句柄##

file 類直接支持上下文管理器 API，但一些其他有打開句柄的對象並不具備這個功能。contextlib 標准庫文檔中給了一個關閉從 urllib.urlopen() 返回的對象的例子。還有許多有 close() 方法但不支持上下文管理器 API 的類。為了確保句柄被關閉，可以使用 closing() 來給它創建一個上下文管理器。

contextlib_closing.py

import contextlib class Door: def __init__(self): print(' __init__()') self.status = 'open' def close(self): print(' close()') self.status = 'closed' print('Normal Example:') with contextlib.closing(Door()) as door: print(' inside with statement: {}'.format(door.status)) print(' outside with statement: {}'.format(door.status)) print('\nError handling example:') try: with contextlib.closing(Door()) as door: print(' raising from inside with statement') raise RuntimeError('error message') except Exception as err: print(' Had an error:', err)

這樣，不管 with 代碼塊中會不會有錯誤拋出，句柄總會被關閉。

$ python3 contextlib_closing.py Normal Example: __init__() inside with statement: open close() outside with statement: closed Error handling example: __init__() raising from inside with statement close() Had an error: error message

Cyrbuzz 翻譯於

由 Summer 審閱

忽略異常##

我們經常需要忽略拋出的異常，因為這樣的異常表示期望的狀態已經達到了，或者它可以是被忽略的異常。常用的做法是寫 try:except 語句然后在 except 里只寫一句 pass。

contextlib_ignore_error.py

import contextlib class NonFatalError(Exception): pass def non_idempotent_operation(): raise NonFatalError( 'The operation failed because of existing state' ) try: print('trying non-idempotent operation') non_idempotent_operation() print('succeeded!') except NonFatalError: pass print('done')

比如這樣，拋出異常然后被忽略。

$ python3 contextlib_ignore_error.py trying non-idempotent operation done

try:except 形式的忽略可以被 contextlib.suppress() 來代替以更加顯式的處理發生在 with 代碼塊中的異常類。

contextlib_suppress.py

import contextlib class NonFatalError(Exception): pass def non_idempotent_operation(): raise NonFatalError( 'The operation failed because of existing state' ) with contextlib.suppress(NonFatalError): print('trying non-idempotent operation') non_idempotent_operation() print('succeeded!') print('done')

更新后的版本，異常也被完全丟棄了。

$ python3 contextlib_suppress.py trying non-idempotent operation done

Cyrbuzz 翻譯於

由 Summer 審閱

重定向輸出流##

設計得不好的庫代碼中可能直接寫了 sys.stdout 或 sys.stderr 這樣的語句，沒有提供參數來配置不同的輸出路口。

redirect_stdout() 和 redirect_stderr() 上下文管理器可以用於捕獲沒有提供接受新的輸出參數的函數中的輸出。

contextlib_redirect.py

from contextlib import redirect_stdout, redirect_stderr import io import sys def misbehaving_function(a): sys.stdout.write('(stdout) A: {!r}\n'.format(a)) sys.stderr.write('(stderr) A: {!r}\n'.format(a)) capture = io.StringIO() with redirect_stdout(capture), redirect_stderr(capture): misbehaving_function(5) print(capture.getvalue())

本例中， misbehaving_function() 同時寫了 stdout 和 stderr ，不過后面兩個上下文管理器都使用了同一個 io.StringIO實例將其捕獲用於之后的使用。
In this example, misbehaving_function() writes to both stdout and stderr, but the two context managers send that output to the same io.StringIO instance where it is saved to be used later.

$ python3 contextlib_redirect.py (stdout) A: 5 (stderr) A: 5

注意##

redirect_stdout() 和 redirect_stderr() 會通過替換 sys 模塊中的對象來修改全局的輸出流，請小心使用。而且該函數不是線程安全的，可能會擾亂輸出到終端上的其他操作的標准輸出。

動態上下文管理器棧##

大多數上下文管理器一次只會操作一個對象，比如單個文件或單個數據庫句柄。這些情況中對象都是提前知道的，使用上下文管理器也都可以圍繞這個對象展開。不過在另一些情況中，可能需要創建一個未知數量的上下文，同時希望控制流退出上下文時這些上下文管理器也全部執行清理功能。 ExitStack 就是用來處理這些動態情況的。

ExitStack 實例維護一個包含清理回調的棧。這些回調都會被放在上下文中，任何被注冊的回調都會在控制流退出上下文時以倒序方式被調用。這有點像嵌套了多層的 with 語句，除了它們是被動態創建的。

上下文管理器棧##

有幾種填充 ExitStack 的方式。本例使用 enter_context() 來將一個新的上下文管理器添加入棧。

contextlib_exitstack_enter_context.py

import contextlib @contextlib.contextmanager def make_context(i): print('{} entering'.format(i)) yield {} print('{} exiting'.format(i)) def variable_stack(n, msg): with contextlib.ExitStack() as stack: for i in range(n): stack.enter_context(make_context(i)) print(msg) variable_stack(2, 'inside context')

enter_context() 首先會調用上下文管理器中的__enter__() 方法，然后把它的 __exit__() 注冊為一個回調以便讓棧調用。

$ python3 contextlib_exitstack_enter_context.py 0 entering 1 entering inside context 1 exiting 0 exiting

ExitStack 中的上下文管理器會像一系列嵌套的 with 一樣。任何發生在上下文中的錯誤都會交給上下文管理器的正常錯誤處理系統去處理。下面的上下文管理器類們可以說明傳遞方式。

contextlib_context_managers.py

import contextlib class Tracker: "用於提醒上下文信息的基礎類" def __init__(self, i): self.i = i def msg(self, s): print(' {}({}): {}'.format( self.__class__.__name__, self.i, s)) def __enter__(self): self.msg('entering') class HandleError(Tracker): "處理任何接收到的異常." def __exit__(self, *exc_details): received_exc = exc_details[1] is not None if received_exc: self.msg('handling exception {!r}'.format( exc_details[1])) self.msg('exiting {}'.format(received_exc)) # 返回布爾類型的值代表是否已經處理了該異常。 return received_exc class PassError(Tracker): "傳遞任何接收到的異常。" def __exit__(self, *exc_details): received_exc = exc_details[1] is not None if received_exc: self.msg('passing exception {!r}'.format( exc_details[1])) self.msg('exiting') # 返回False，表示沒有處理這個異常。 return False class ErrorOnExit(Tracker): "拋出個異常" def __exit__(self, *exc_details): self.msg('throwing error') raise RuntimeError('from {}'.format(self.i)) class ErrorOnEnter(Tracker): "拋出個異常." def __enter__(self): self.msg('throwing error on enter') raise RuntimeError('from {}'.format(self.i)) def __exit__(self, *exc_info): self.msg('exiting')

例子中的類會被包含在 variable_stack() 中使用（見上面的代碼），variable_stack() 把上下文管理器放到 ExitStack 中使用，逐一建立起上下文。下面的例子我們將傳遞不同的上下文管理器來測試錯誤處理結果。首先我們測試無異常的常規情況。

print('No errors:') variable_stack([ HandleError(1), PassError(2), ])

之后，我們做一個在棧末的處理異常的例子，這樣的話所有已經打開的上下文管理器會隨着棧的釋放而關閉。

print('\nError at the end of the context stack:') variable_stack([ HandleError(1), HandleError(2), ErrorOnExit(3), ])

接着，我們做一個在棧中間處理異常的例子，這時我們會看到發生錯誤時某些上下文已經關閉，所以那些上下文不會受到這個異常的影響。

print('\nError in the middle of the context stack:') variable_stack([ HandleError(1), PassError(2), ErrorOnExit(3), HandleError(4), ])

最后，放一個不處理的異常，然后傳到上層調用它的代碼中。

try: print('\nError ignored:') variable_stack([ PassError(1), ErrorOnExit(2), ]) except RuntimeError: print('error handled outside of context')

我們可以看到，如果棧中的任何一個上下文管理器接收到這個異常然后返回了一個 True 的話，這個異常就會就此消失，不會再進行傳播，否則就會一直傳遞下去。

$ python3 contextlib_exitstack_enter_context_errors.py No errors: HandleError(1): entering PassError(2): entering PassError(2): exiting HandleError(1): exiting False outside of stack, any errors were handled Error at the end of the context stack: HandleError(1): entering HandleError(2): entering ErrorOnExit(3): entering ErrorOnExit(3): throwing error HandleError(2): handling exception RuntimeError('from 3',) HandleError(2): exiting True HandleError(1): exiting False outside of stack, any errors were handled Error in the middle of the context stack: HandleError(1): entering PassError(2): entering ErrorOnExit(3): entering HandleError(4): entering HandleError(4): exiting False ErrorOnExit(3): throwing error PassError(2): passing exception RuntimeError('from 3',) PassError(2): exiting HandleError(1): handling exception RuntimeError('from 3',) HandleError(1): exiting True outside of stack, any errors were handled Error ignored: PassError(1): entering ErrorOnExit(2): entering ErrorOnExit(2): throwing error PassError(1): passing exception RuntimeError('from 2',) PassError(1): exiting error handled outside of context

Cyrbuzz 翻譯於

由 Summer 審閱

任意上下文回調##

ExitStack 也支持關閉上下文時有其他回調，使用這種方法無需經由上下文管理器控制，可以更方便得清理資源。

contextlib_exitstack_callbacks.py

import contextlib def callback(*args, **kwds): print('closing callback({}, {})'.format(args, kwds)) with contextlib.ExitStack() as stack: stack.callback(callback, 'arg1', 'arg2') stack.callback(callback, arg3='val3')

相當於所有上下文管理器的 __exit__()，這些回調的調用順序也是倒序的。

$ python3 contextlib_exitstack_callbacks.py closing callback((), {'arg3': 'val3'}) closing callback(('arg1', 'arg2'), {})

不管有沒有錯誤發生，這些回調總會被調用，同時也不會對是否發生了錯誤有任何信息。最后這些回調的返回值也不會有任何作用。

contextlib_exitstack_callbacks_error.py

import contextlib def callback(*args, **kwds): print('closing callback({}, {})'.format(args, kwds)) try: with contextlib.ExitStack() as stack: stack.callback(callback, 'arg1', 'arg2') stack.callback(callback, arg3='val3') raise RuntimeError('thrown error') except RuntimeError as err: print('ERROR: {}'.format(err))

也正因為這些回調無法訪問到錯誤，所以也就無法通過在上下文管理器棧中傳遞錯誤來忽略它。

$ python3 contextlib_exitstack_callbacks_error.py closing callback((), {'arg3': 'val3'}) closing callback(('arg1', 'arg2'), {}) ERROR: thrown error

這樣的回調提供了一種便捷的方式定義清理邏輯而無需創建一個多余的新的上下文管理器類。為了提高可讀性，具體邏輯也可以寫在內聯函數中，callback() 也可以作為裝飾器使用。

contextlib_exitstack_callbacks_decorator.py

import contextlib with contextlib.ExitStack() as stack: @stack.callback def inline_cleanup(): print('inline_cleanup()') print('local_resource = {!r}'.format(local_resource)) local_resource = 'resource created in context' print('within the context')

把callback() 作為裝飾器使用時無法給被注冊的函數指定參數。不過，如果清理函數作為內聯定義，作用域規則也給了它訪問調用它的代碼中變量的權力。

$ python3 contextlib_exitstack_callbacks_decorator.py within the context inline_cleanup() local_resource = 'resource created in context'

Cyrbuzz 翻譯於

由 Summer 審閱

局部棧##

有時我們需要創建一個復雜的上下文時，如果上下文無法完全構造出來，使用局部棧可以有效打斷某一操作。不過如果設置正確的話，一段時間之后也會清理其中所有的資源。舉個例子，在單個上下文中，如果某一操作需要多個長時間存活的網絡連接，其中某一連接失效時，最好的情況是不進行這個操作。但如果所有連接都正確打開，那也需要它保持正常操作。 ExitStack 中的 pop_all() 則適用於這種情況。

pop_all() 會在被調用時清理棧中所有的上下文管理器和回調，並返回一個包含與之前的棧相同內容的新棧。原棧完成操作后，可以新棧的 close() 方法清理所有資源。

contextlib_exitstack_pop_all.py

import contextlib from contextlib_context_managers import * def variable_stack(contexts): with contextlib.ExitStack() as stack: for c in contexts: stack.enter_context(c) # 返回新棧的 close() 方法作為清理函數使用。 return stack.pop_all().close # 直接返回None，表示 ExitStack 沒有完成干凈的初始化 # 它的清理過程已經發生。 return None print('No errors:') cleaner = variable_stack([ HandleError(1), HandleError(2), ]) cleaner() print('\nHandled error building context manager stack:') try: cleaner = variable_stack([ HandleError(1), ErrorOnEnter(2), ]) except RuntimeError as err: print('caught error {}'.format(err)) else: if cleaner is not None: cleaner() else: print('no cleaner returned') print('\nUnhandled error building context manager stack:') try: cleaner = variable_stack([ PassError(1), ErrorOnEnter(2), ]) except RuntimeError as err: print('caught error {}'.format(err)) else: if cleaner is not None: cleaner() else: print('no cleaner returned')

繼續使用之前定義好的上下文管理器類，不一樣的是 ErrorOnEnter 會在 __enter__() 產生錯誤而不是在 __exit__() 中。 variable_stack() 內部的邏輯是如果所有的上下文都成功進入且無錯誤產生，則會返回新ExitStack 的 close() 方法。如果處理了一個錯誤，則返回 None 指代清理工作已經完成了。如果發生錯誤但並未處理，則清理局部棧，之后錯誤會繼續傳遞。

$ python3 contextlib_exitstack_pop_all.py No errors: HandleError(1): entering HandleError(2): entering HandleError(2): exiting False HandleError(1): exiting False Handled error building context manager stack: HandleError(1): entering ErrorOnEnter(2): throwing error on enter HandleError(1): handling exception RuntimeError('from 2',) HandleError(1): exiting True no cleaner returned Unhandled error building context manager stack: PassError(1): entering ErrorOnEnter(2): throwing error on enter PassError(1): passing exception RuntimeError('from 2',) PassError(1): exiting caught error from 2