Python的上下文管理器

发表于 2017年7月12日 11:19 更新于 2020年10月12日 9:50 分类于 Python 阅读次数：

Python中的上下文管理器可以允许精确地分配和释放资源，最常用的就是使用with语句，比如：

with open('/tmp/file_x', 'w') as file_x:
    file_x.write('Hello')

当with结束，文件也会被安全的关闭。不用担心回收资源的问题了。

如果一个自定义的类也想支持类似的调用方式，需要实现__enter__(self)和__exit__(self, type, value, traceback)这两个方法，具体的：

class File(object):
    def __init__(self, file_name, method):
        self.file_obj = open(file_name, method)
    def __enter__(self):
        return self.file_obj
    def __exit__(self, type, value, traceback):
        self.file_obj.close()

其中__enter__方法将打开的文件返回给with语句。
对于__exit__(self, type, value, traceback)方法，会在with语句退出时调用，如果在执行中发现异常，则异常的type,value和traceback会被传递给__exit__方法，在__exit__中可以对异常进行相应的处理，如果最终
__exit__方法返回None，则认为异常被正确处理了，如果返回的不是None，则这个异常会被with抛出，期待上层进行相应的处理。

除了上面的方法，Python还提供了一个contextlib模块，通过这个模块加上装饰器(decorators)和生成器(generators)，也能实现类似的功能：

from contextlib import contextmanager

@contextmanager
def open_file(name):
    f = open(name, 'w')
    yield f
    f.close()

这样在使用中，open_file变成了一个生成器，所以contextmanager可以通过调用这个生成器next()方法控制资源的释放，具体的源代码在这里：

# 代码有所省略，具体可以参考： https://github.com/python/cpython/blob/master/Lib/contextlib.py
class _GeneratorContextManager(_GeneratorContextManagerBase,
                               AbstractContextManager,
                               ContextDecorator):
    """Helper for @contextmanager decorator."""

    def _recreate_cm(self):
        return self.__class__(self.func, self.args, self.kwds)

    def __enter__(self):
        try:
            return next(self.gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, type, value, traceback):
        if type is None:
            try:
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:
                value = type()
            try:
                self.gen.throw(type, value, traceback)
            except StopIteration as exc:
                return exc is not value
            except RuntimeError as exc:
                if exc is value:
                    return False
                if type is StopIteration and exc.__cause__ is value:
                    return False
                raise
            except:
                if sys.exc_info()[1] is value:
                    return False
                raise
            raise RuntimeError("generator didn't stop after throw()")

def contextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper

参考：