每週一個 Python 模塊 | contextlib

時間 2019-11-30

標籤每週一個 python 模塊 contextlib 欄目 Python 简体版

原文原文鏈接

用於建立和使用上下文管理器的實用程序。python

contextlib 模塊包含用於處理上下文管理器和 with 語句的實用程序。git

Context Manager API

上下文管理器負責一個代碼塊內的資源，從進入塊時建立到退出塊後清理。例如，文件上下文管理器 API，在完成全部讀取或寫入後來確保它們已關閉。github

with open('/tmp/pymotw.txt', 'wt') as f:
    f.write('contents go here')
# file is automatically closed
複製代碼

with 語句啓用了上下文管理器，API 涉及兩種方法：當執行流進入內部代碼塊時運行 __enter__() 方法，它返回要在上下文中使用的對象。當執行流離開 with 塊時，調用上下文管理器的 __exit__() 方法來清理正在使用的任何資源。數據庫

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')
    
# output
# __init__()
# __enter__()
# Doing work in the context
# __exit__()
複製代碼

組合上下文管理器和 with 語句是一種更簡潔的 try:finally 塊，即便引起了異常，也老是調用上下文管理器的 __exit__() 方法。api

__enter__() 方法能夠返回與 as 子句中指定的名稱關聯的任何對象。在此示例中，Context 返回使用打開上下文的對象。安全

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()
    
# output
# Context.__init__()
# Context.__enter__()
# WithinContext.__init__(<__main__.Context object at 0x101f046d8>)
# WithinContext.do_something()
# Context.__exit__()
# WithinContext.__del__
複製代碼

與變量關聯的值 c 是返回的 __enter__() 對象，該對象不必定是 Context 在 with 語句中建立的實例。bash

__exit__() 方法接收包含 with 塊中引起的任何異常的詳細信息的參數。網絡

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')
    
# output
# __init__(True)
# __enter__()
# __exit__()
# exc_type = <class 'RuntimeError'>
# exc_val = error message handled
# exc_tb = <traceback object at 0x101c94948>
# 
# __init__(False)
# __enter__()
# __exit__()
# exc_type = <class 'RuntimeError'>
# exc_val = error message propagated
# exc_tb = <traceback object at 0x101c94948>
# Traceback (most recent call last):
# File "contextlib_api_error.py", line 34, in <module>
# raise RuntimeError('error message propagated')
# RuntimeError: error message propagated
複製代碼

若是上下文管理器能夠處理異常，__exit__() 則應返回 true 值以指示不須要傳播該異常，返回 false 會致使在 __exit__() 返回後從新引起異常。數據結構

做爲函數裝飾器的上下文管理器

類 ContextDecorator 增長了對常規上下文管理器類的支持，使它們能夠像用上下文管理器同樣用函數裝飾器。

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')

# output
# __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__() 返回的值在被裝飾的函數內部不可用，這與使用 with 和 as 時不一樣，傳遞給裝飾函數的參數以一般方式提供。

從生成器到上下文管理器

經過用 __enter__() 和 __exit__() 方法編寫類來建立上下文管理器的傳統方式並不困難。可是，有時候徹底寫出全部內容對於一些微不足道的上下文來講是沒有必要的。在這些狀況下，使用 contextmanager() 裝飾器將生成器函數轉換爲上下文管理器。

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')
    
# output
# 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
複製代碼

生成器應該初始化上下文，只產生一次，而後清理上下文。若是有的話，產生的值綁定到 as 子句中的變量。with 塊內的異常在生成器內從新引起，所以能夠在那裏處理它們。

contextmanager() 返回的上下文管理器派生自 ContextDecorator，所以它也能夠做爲函數裝飾器使用。

@contextlib.contextmanager
def make_context():
    print(' entering')
    try:
        # Yield control, but not a value, because any value
        # yielded is not available when the context manager
        # is used as a decorator.
        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'))

# output
# 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.7/contextlib.py", line 74, in inner
# return func(*args, **kwds)
# File "contextlib_contextmanager_decorator.py", line 33, in
# throw_error
# raise err
# ValueError: this exception is not handled
複製代碼

如上例所示，當上下文管理器用做裝飾器時，生成器產生的值在被裝飾的函數內不可用，傳遞給裝飾函數的參數仍然可用，如 throw_error() 中所示。

關閉打開句柄

file 類支持上下文管理器 API，但表明打開句柄的一些其餘對象並不支持。標準庫文檔中給出的 contextlib 示例是 urllib.urlopen() 返回的對象。還有其餘遺留類使用 close() 方法，但不支持上下文管理器 API。要確保句柄已關閉，請使用 closing() 爲其建立上下文管理器。

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)
    
# output
# 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
複製代碼

不管 with 塊中是否有錯誤，句柄都會關閉。

忽略異常

忽略異常的最多見方法是使用語句塊 try:except，而後在語句 except 中只有 pass。

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')

# output
# trying non-idempotent operation
# done
複製代碼

在這種狀況下，操做失敗並忽略錯誤。

try:except 能夠被替換爲 contextlib.suppress()，更明確地抑制類異常在 with 塊的任何地方發生。

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')

# output
# trying non-idempotent operation
# done
複製代碼

在此更新版本中，異常將徹底丟棄。

重定向輸出流

設計不良的庫代碼可能直接寫入 sys.stdout 或 sys.stderr，不提供參數來配置不一樣的輸出目的地。若是源不能被改變接受新的輸出參數時，可使用 redirect_stdout() 和 redirect_stderr() 上下文管理器捕獲輸出。

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())

# output
# (stdout) A: 5
# (stderr) A: 5
複製代碼

在此示例中，misbehaving_function() 寫入 stdout 和 stderr，但兩個上下文管理器將該輸出發送到同一 io.StringIO，保存它以便稍後使用。

注意：redirect_stdout() 和 redirect_stderr() 經過替換 sys 模塊中的對象來修改全局狀態，應當心使用。這些函數不是線程安全的，而且可能會干擾指望將標準輸出流附加到終端設備的其餘操做。

動態上下文管理器堆棧

大多數上下文管理器一次操做一個對象，例如單個文件或數據庫句柄。在這些狀況下，對象是事先已知的，而且使用上下文管理器的代碼能夠圍繞該對象構建。在其餘狀況下，程序可能須要在上下文中建立未知數量的對象，同時但願在控制流退出上下文時清除全部對象。ExitStack 函數就是爲了處理這些更動態的狀況。

ExitStack 實例維護清理回調的堆棧數據結構。回調在上下文中顯式填充，而且當控制流退出上下文時，以相反的順序調用已註冊的回調。就像有多個嵌套 with 語句，只是它們是動態創建的。

堆疊上下文管理器

有幾種方法能夠填充 ExitStack。此示例用於 enter_context() 向堆棧添加新的上下文管理器。

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')

# output
# 0 entering
# 1 entering
# inside context
# 1 exiting
# 0 exiting
複製代碼

enter_context() 首先調用 __enter__() 上下文管理器，而後將 __exit__() 方法註冊爲在棧撤消時調用的回調。

上下文管理器 ExitStack 被視爲處於一系列嵌套 with 語句中。在上下文中的任何位置發生的錯誤都會經過上下文管理器的正常錯誤處理進行傳播。這些上下文管理器類說明了錯誤傳播的方式。

# contextlib_context_managers.py 
import contextlib


class Tracker:
    "Base class for noisy context managers."

    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):
    "If an exception is received, treat it as handled."

    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 Boolean value indicating whether the exception
        # was handled.
        return received_exc


class PassError(Tracker):
    "If an exception is received, propagate it."

    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')
        # Return False, indicating any exception was not handled.
        return False


class ErrorOnExit(Tracker):
    "Cause an exception."

    def __exit__(self, *exc_details):
        self.msg('throwing error')
        raise RuntimeError('from {}'.format(self.i))


class ErrorOnEnter(Tracker):
    "Cause an exception."

    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()，它使用上下文管理器來構造 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
複製代碼

任意上下文回調

ExitStack 還支持關閉上下文的任意回調，從而能夠輕鬆清理不經過上下文管理器控制的資源。

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')
    
# output
# closing callback((), {'arg3': 'val3'})
# closing callback(('arg1', 'arg2'), {})
複製代碼

與 __exit__() 完整上下文管理器的方法同樣，回調的調用順序與它們的註冊順序相反。

不管是否發生錯誤，都會調用回調，而且不會給出有關是否發生錯誤的任何信息。它們的返回值被忽略。

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))
    
# output
# closing callback((), {'arg3': 'val3'})
# closing callback(('arg1', 'arg2'), {})
# ERROR: thrown error
複製代碼

由於它們沒法訪問錯誤，因此回調沒法經過其他的上下文管理器堆棧阻止異常傳播。

回調能夠方便清楚地定義清理邏輯，而無需建立新的上下文管理器類。爲了提升代碼可讀性，該邏輯能夠封裝在內聯函數中，callback() 能夠用做裝飾器。

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')
    
# output
# within the context
# inline_cleanup()
# local_resource = 'resource created in context'
複製代碼

沒法爲使用裝飾器形式註冊的 callback() 函數指定參數。可是，若是清理回調是內聯定義的，則範圍規則容許它訪問調用代碼中定義的變量。

部分堆棧

有時，在構建複雜的上下文時，若是上下文沒法徹底構建，能夠停止操做，可是若是延遲清除全部資源，則可以正確設置全部資源。例如，若是操做須要多個長期網絡鏈接，則最好不要在一個鏈接失敗時啓動操做。可是，若是能夠打開全部鏈接，則須要保持打開的時間長於單個上下文管理器的持續時間。能夠在此方案中使用 ExitStack 的 pop_all() 方法。

pop_all() 從調用它的堆棧中清除全部上下文管理器和回調，並返回一個預先填充了相同上下文管理器和回調的新堆棧。在原始堆棧消失以後，能夠稍後調用新堆棧的 close() 方法來清理資源。

import contextlib

from contextlib_context_managers import *


def variable_stack(contexts):
    with contextlib.ExitStack() as stack:
        for c in contexts:
            stack.enter_context(c)
        # Return the close() method of a new stack as a clean-up
        # function.
        return stack.pop_all().close
    # Explicitly return None, indicating that the ExitStack could
    # not be initialized cleanly but that cleanup has already
    # occurred.
    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')
        
# output
# 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
複製代碼

此示例使用前面定義的相同上下文管理器類，其差別是 ErrorOnEnter 產生的錯誤是 __enter__() 而不是 __exit__()。在 variable_stack() 內，若是輸入的全部上下文都沒有錯誤，則返回一個 ExitStack 的 close() 方法。若是發生處理錯誤，則 variable_stack() 返回 None 來表示已完成清理工做。若是發生未處理的錯誤，則清除部分堆棧並傳播錯誤。

相關文檔：

https://pymotw.com/3/contextlib/index.html