pythonic context manager知多少

Context Managers 是我最喜欢的 python feature 之一，在恰当的时机使用 context manager 使代码更加简洁、清晰，更加安全，复用性更好，更加 pythonic。本文简单介绍一下其使用方法以及常见使用场景。

本文地址：https://www.cnblogs.com/xybaby/p/13202496.html

with statement and context manager

Python’s with statement supports the concept of a runtime context defined by a context manager

new statement "with" to the Python language to make it possible to factor out standard uses of try/finally statements.

在 pep0343 中，通过引入 context manager protocol 来支持 With statement ， context manager 是用来管理 context（上下文）的，即保证程序要保持一种特定的状态 -- 无论是否发生异常。可以说，context manager 简化了对 try-finally 的使用，而且更加安全，更加便于使用。

Transforming Code into Beautiful, Idiomatic Python 中，指出了 context manager 的最显著的优点：

• Helps separate business logic from administrative logic
• Clean, beautiful tools for factoring code and improving code reuse

最广为人知的例子，就是通过 with statement 来读写文件，代码如下：

with open('test.txt') as f:
    contect = f.read()
    handle_content(content)

上面的代码几乎等价于

f = open('test.txt') 
try:
    contect = f.read()
    handle_content(content)
finally:
    f.close()

注意，上面的finally的作用就是保证file.close一定会被调用，也就是资源一定会释放。不过，很多时候，都会忘了去写这个finally，而 with statement 就彻底避免了这个问题。

从上述两段代码也可以看出，with statement 更加简洁，而且将核心的业务逻辑（从文件中读取、处理数据）与其他逻辑（打开、关系文件）相分离，可读性更强。

实现context manager protocol

一个类只要定义了__enter__、__exit__方法就实现了context manager 协议

object.__enter__(self)
Enter the runtime context related to this object. The with statement will bind this method’s return value to the target(s) specified in the as clause of the statement, if any.

object.__exit__(self, exc_type, exc_value, traceback)
Exit the runtime context related to this object. The parameters describe the exception that caused the context to be exited. If the context was exited without an exception, all three arguments will be None.

If an exception is supplied, and the method wishes to suppress the exception (i.e., prevent it from being propagated), it should return a true value. Otherwise, the exception will be processed normally upon exit from this method.

Note that __exit__() methods should not reraise the passed-in exception; this is the caller’s responsibility.

__enter__方法在进入这个 context 的时候调用，返回值赋值给 with as X 中的 X

__exit__方法在退出 context 的时候调用，如果没有异常，后三个参数为 None。如果返回值为 True，则Suppress Exception，所以除非特殊情况都应返回 False。另外注意， __exit__方法本身不应该抛出异常。

例子：BlockGuard

在看c++代码（如mongodb源码）的时候，经常看见其用 RAII 实现BlockGuard, 用以保证在离开 Block 的时候执行某些动作，同时，也提供手段来取消执行。

下面用python实现一下：

class BlockGuard(object):
	def __init__(self, fn, *args, **kwargs):
		self._fn = fn
		self._args = args
		self._kwargs = kwargs
		self._canceled = False

	def __enter__(self):
		return self

	def __exit__(self, exc_type, exc_value, traceback):
		if not self._canceled:
			self._fn(*self._args, **self._kwargs)
		self._fn = None
		self._args = None
		self._kwargs = None
		return False

	def cancel(self):
		self._canceled = True


def foo():
	print 'sth should be called'


def test_BlockGuard(cancel_guard):
	print 'test_BlockGuard'
	with BlockGuard(foo) as guard:
		if cancel_guard:
			guard.cancel()
	print 'test_BlockGuard  finish'

用yield实现context manager

标准库 contextlib 中提供了一些方法，能够简化我们使用 context manager，如 contextlib.contextmanager(func) 使我们
无需再去实现一个包含__enter__ __exit__方法的类。

The function being decorated must return a generator-iterator when called. This iterator must yield exactly one value, which will be bound to the targets in the with statement’s as clause, if any.

例子如下：

from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwds):
    # Code to acquire resource, e.g.:
    resource = acquire_resource(*args, **kwds)
    try:
        yield resource
    finally:
        # Code to release resource, e.g.:
        release_resource(resource)

>>> with managed_resource(timeout=3600) as resource:
...     # Resource is released at the end of this block,
...     # even if code in the block raises an exception

需要注意的是：

• 一定要写 try finally，才能保证release_resource逻辑一定被调用
• 除非特殊情况，不再 catch exception，这就跟 __exit__ 一般不返回True一样

例子: no_throw

这是业务开发中的一个需求, 比如观察者模式，不希望因为其中一个观察者出了 trace 就影响后续的观察者，就可以这样做：

from contextlib import contextmanager

@contextmanager
def no_throw(*exceptions):
	try:
		yield
	except exceptions:
		pass

def notify_observers(seq):
	for fn in [sum, len, max, min]:
		with no_throw(Exception):
			print "%s result %s" % (fn.__name__, fn(seq))

if __name__ == '__main__':
	notify_observers([])

在python 3.x 的 contexlib 中，就提供了一个contextlib.suppress(*exceptions), 实现了同样的效果。

context manager 应用场景

context manager 诞生的初衷就在于简化 try-finally，因此就适合应用于在需要 finally 的地方，也就是需要清理的地方，比如

• 保证资源的安全释放，如 file、lock、semaphore、network connection 等
• 临时操作的复原，如果一段逻辑有 setup、prepare，那么就会对应 cleanup、teardown。

对于第一种情况，网络连接释放的例子，后面会结合 pymongo 的代码展示。

在这里先来看看第二种用途：保证代码在一个临时的、特殊的上下文（context）中执行，且在执行结束之后恢复到之前的上下文环境。

改变工作目录

from contextlib import contextmanager
import os

@contextmanager
def working_directory(path):
    current_dir = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(current_dir)

with working_directory("data/stuff"):
    pass

临时文件、文件夹

很多时候会产生一堆临时文件，比如build的中间状态，这些临时文件都需要在结束之后清除。

from tempfile import mkdtemp
from shutil import rmtree

@contextmanager
def temporary_dir(*args, **kwds):
    name = mkdtemp(*args, **kwds)
    try:
        yield name
    finally:
        shutil.rmtree(name)

with temporary_dir() as dirname:
    pass

重定向标准输出、标准错误

@contextmanager
def redirect_stdout(fileobj):
    oldstdout = sys.stdout
    sys.stdout = fileobj
    try:
        yield fieldobj
    finally:
        sys.stdout = oldstdout

在 python3.x 中，已经提供了 contextlib.redirect_stdout contextlib.redirect_stderr 实现上述功能

调整logging level

这个在查问题的适合非常有用，一般生产环境不会输出 debug level 的日志，但如果出了问题，可以临时对某些制定的函数调用输出debug 日志

from contextlib import contextmanager
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

ch = logging.StreamHandler()
ch.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(ch)


@contextmanager
def change_log_level(level):
	old_level = logger.getEffectiveLevel()
	try:
		logger.setLevel(level)
		yield
	finally:
		logger.setLevel(old_level)


def test_logging():
	logger.debug("this is a debug message")
	logger.info("this is a info message")
	logger.warn("this is a warning message")

with change_log_level(logging.DEBUG):
	test_logging()

pymongo中的context manager使用

在 pymongo 中，封装了好几个 context manager，用以

• 管理 semaphore
• 管理 connection
• 资源清理

而且，在 pymongo 中，给出了嵌套使用 context manager 的好例子，用来保证 socket 在使用完之后一定返回连接池(pool)。

# server.py
@contextlib.contextmanager
def get_socket(self, all_credentials, checkout=False):
    with self.pool.get_socket(all_credentials, checkout) as sock_info:
        yield sock_info
        
# pool.py
@contextlib.contextmanager
def get_socket(self, all_credentials, checkout=False):
    sock_info = self._get_socket_no_auth()
    try:
        sock_info.check_auth(all_credentials)
        yield sock_info
    except:
        # Exception in caller. Decrement semaphore.
        self.return_socket(sock_info)
        raise
    else:
        if not checkout:
            self.return_socket(sock_info)

可以看到，server.get_socket 调用了 pool.get_socket, 使用 server.get_socket 的代码完全不了解、也完全不用关心 socket 的释放细节，如果把 try-except-finally-else 的逻辑移到所有使用socket的地方，代码就会很丑、很臃肿。

比如，在mongo_client 中需要使用到 socket：

with server.get_socket(all_credentials) as sock_info:
    sock_info.authenticate(credentials)

references

With statement

Context Managers

contextlib

what-is-the-python-with-statement-designed-for

Transforming Code into Beautiful, Idiomatic Python