4.5. contextlib — 上下文管理器工具

Python 3 标准库实例教程 /

本节目标: 创建和使用基于上下文管理器的工具

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

函数装饰器方式的上下文管理器

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

关闭打开的句柄

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

忽略异常

我们经常需要忽略抛出的异常，因为这样的异常表示期望的状态已经达到了，或者它可以是被忽略的异常。常用的做法是写 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

重定向输出流

设计得不好的库代码中可能直接写了 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

任意上下文回调

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'

局部栈

有时我们需要创建一个复杂的上下文时，如果上下文无法完全构造出来，使用局部栈可以有效打断某一操作。不过如果设置正确的话，一段时间之后也会清理其中所有的资源。举个例子，在单个上下文中，如果某一操作需要多个长时间存活的网络连接，其中某一连接失效时，最好的情况是不进行这个操作。但如果所有连接都正确打开，那也需要它保持正常操作。 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

参阅

contextlib 标准库文档

PEP 343 -- with 语句参考

上下文管理器类型参考 -- 上下文管理器 API 标准库文档参考

使用上下文管理器 -- 上下文管理器 API 指南

Python 3.3+ 资源管理器 ExitStack ! -- 使用 ExitStack 编写安全代码 -- Barry Warsaw.

本文章首发在 LearnKu.com 网站上。

本译文仅用于学习和交流目的，转载请务必注明文章译者、出处、和本文链接
我们的翻译工作遵照 CC 协议，如果我们的工作有侵犯到您的权益，请及时联系我们。

原文地址：https://learnku.com/docs/pymotw/contextl...

译文地址：https://learnku.com/docs/pymotw/contextl...

4.5. contextlib — 上下文管理器工具

上下文管理器 API

函数装饰器方式的上下文管理器

从生成器到上下文管理器

关闭打开的句柄

忽略异常

重定向输出流

注意

动态上下文管理器栈

上下文管理器栈

任意上下文回调

局部栈

参阅

社区赞助商

关于 LearnKu

资源推荐

服务提供商

其他信息

4.5. contextlib — 上下文管理器工具

上下文管理器 API

函数装饰器方式的上下文管理器

从生成器到上下文管理器

关闭打开的句柄

忽略异常

重定向输出流

注意

动态上下文管理器栈

上下文管理器栈

任意上下文回调

局部栈

参阅

社区赞助商

关于 LearnKu

资源推荐

服务提供商

其他信息

请登录