每週一個 Python 模組 | contextlib

yongxinz發表於2019-03-03
PythonContext

專欄地址：每週一個 Python 模組

同時，也歡迎關注我的微信公眾號 AlwaysBeta，更多精彩內容等你來。

每週一個 Python 模組 | contextlib

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

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

Context Manager API

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

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__() 方法。

__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 語句中建立的例項。

__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 來表示已完成清理工作。如果發生未處理的錯誤，則清除部分堆疊並傳播錯誤。

相關文件：

pymotw.com/3/contextli…

相關文章