使用@contextmanager装饰器实现上下文管理器

一般来讲，实现上下文管理器，须要编写一个带有__enter__和 __exit__的类，相似这样：

class ListTransaction:

    def __init__(self, orig_list):
        self.orig_list = orig_list
        self.working = list(orig_list)

    def __enter__(self):
        return self.working

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.orig_list[:] = self.working

然而，在contextlib模块中，还提供了@contextmanager装饰器，将一个生成器函数当成上下文管理器使用，上面的代码在大部分，是与下面的代码等效的。

本文的list_transaction函数的代码来自：《Python Cookbook》 9.22 以简单的方式定义上下文管理器

from contextlib import contextmanager

@contextmanager
def list_transaction(orig_list):
    working = list(orig_list)
    yield working
    orig_list[:] = working

先逐一分析上面的代码：

1. 由于list是可变类型，因此经过list(orig_list)，对值进行复制，建立一个新的list，即working。
2. 以yield为分隔，在yield以前的代码，包括yield working，会在contextmanager装饰器的__enter__方法中被调用
3. 代码在执行到yield时暂停，同时yield working，会将working产出。yield产出的值，做为__enter__的返回值，赋值给as以后的变量
4. 当with块的代码执行完成后， 上下文管理器会在yield处恢复，继续执行yield以后的代码。
5. yield 以后的代码，则在contextmanager装饰器中的__exit__方法中被调用

测试代码以下：

当执行过程当中，没有引起异常时，执行正常，输出 [1, 2, 3, 4, 5]

    items_1 = [1, 2, 3]
    with list_transaction(items_1) as working_1:
        working_1.append(4)
        working_1.append(5)
    print(items_1)

当执行过程当中，引起异常时，yield后的代码不会执行，orig_list不会被修改。从而实现事务的效果，orig_list还是 [1, 2, 3]

    items_2 = [1, 2, 3]
    try:
        with list_transaction(items_2) as working_2:
            working_2.append(4)
            working_2.append(5)
            raise RuntimeError('oops')
    except Exception as ex:
        print(ex)
    finally:
        print(items_2)

上下文管理器类与@contextmanager中最大的区别在于对异常的处理。

分析contextmanager的源码可知，@contextmanager装饰器的本质是实例化一个_GeneratorContextManager对象。

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

进一步查看_GeneratorContextManager源码，可知_GeneratorContextManager实现的是一个上下文管理器对象

class _GeneratorContextManager(ContextDecorator):
    """Helper for @contextmanager decorator."""

    def __init__(self, func, args, kwds):
        self.gen = func(*args, **kwds)
        self.func, self.args, self.kwds = func, args, kwds
        # Issue 19330: ensure context manager instances have good docstrings
        doc = getattr(func, "__doc__", None)
        if doc is None:
            doc = type(self).__doc__
        self.__doc__ = doc
        # Unfortunately, this still doesn't provide good help output when
        # inspecting the created context manager instances, since pydoc
        # currently bypasses the instance docstring and shows the docstring
        # for the class instead.
        # See http://bugs.python.org/issue19404 for more details.

    def _recreate_cm(self):
        # _GCM instances are one-shot context managers, so the
        # CM must be recreated each time a decorated function is
        # called
        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
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:
                # Need to force instantiation so we can reliably
                # tell if we get the same exception back
                value = type()
            try:
                self.gen.throw(type, value, traceback)
                raise RuntimeError("generator didn't stop after throw()")
            except StopIteration as exc:
                # Suppress StopIteration *unless* it's the same exception that
                # was passed to throw().  This prevents a StopIteration
                # raised inside the "with" statement from being suppressed.
                return exc is not value
            except RuntimeError as exc:
                # Likewise, avoid suppressing if a StopIteration exception
                # was passed to throw() and later wrapped into a RuntimeError
                # (see PEP 479).
                if exc.__cause__ is value:
                    return False
                raise
            except:
                # only re-raise if it's *not* the exception that was
                # passed to throw(), because __exit__() must not raise
                # an exception unless __exit__() itself failed.  But throw()
                # has to raise the exception to signal propagation, so this
                # fixes the impedance mismatch between the throw() protocol
                # and the __exit__() protocol.
                #
                if sys.exc_info()[1] is not value:
                    raise

简要分析实现的代码：

__enter__方法：

1. self.gen = func(*args, **kwds) 获取生成器函数返回的生成器，并赋值给self.gen
2. with代码块进入__enter__方法时，调用生成器的__next__方法，使代码执行到yield处暂停
3. 将yield产出的值做为__enter__的返回值
4. 由于__enter__方法只会执行一次，若是第一次调用生成器的__next__方法，就抛出StopIteration异常，说明生成器存在问题，则抛出RuntimeError

__exit__方法：

正常执行的状况：

1. def __exit__(self, type, value, traceback)接收三个参数，第一个参数是异常类，第二个参数是异常对象，第三个参数是trackback对象
2. 若是with内的代码执行正常，没有抛出异常，则上述三个参数都为None
3. __exit__代码中首先对type是否None进行判断，若是type为None，说明with代码内部执行正常，因此调用生成器的__next__方法。此时生成器在yield处恢复运行，继续执行yield以后的代码
4. 正常状况下，调用__next__方法，迭代应结束，抛出StopIteration异常；若是没有抛出StopIteration异常，说明生成器存在问题，则抛出RuntimeError

出现异常的状况：

1. 若是type类型不为None，说明在with代码内部执行时出现异常。若是异常对象value为None，则强制使用异常类实例化一个新的异常对象，并赋值给value
2. 使用throw方法，将异常对象value传递给生成器函数，此时生成器在yield处恢复执行，并接收到异常信息
3. 一般状况下，yield语句应该在try except代码块中执行，用于捕获__exit__方法传递给生成器的异常信息，并进行处理
4. 若是生成器函数能够处理异常，迭代完成后，自动抛出StopIteration。
5. __exit__ 捕获并压制StopIteration，除非with内的代码也抛出了StopIteration。return exc is not value，exc是捕获到的StopIteration异常实例，value是with内代码执行时抛出的异常。在__exit__方法中，return True告诉解释器异常已经处理，除此之外，全部的异常都会向上冒泡。
6. 若是生成器没有抛出StopIteration异常，说明迭代没有正常结束，则__exit__方法抛出RuntimeError，一样的，除非with代码块内部也抛出RuntimeError，不然RuntimeError会在__exit__中被捕获而且压制。

 

因此，以类的方式实现的上下文管理器，在引起异常时，__exit__方法内的代码仍会正常执行；

而以生成器函数实现的上下文管理器，在引起异常时，__exit__方法会将异常传递给生成器，若是生成器没法正确处理异常，则yield以后的代码不会执行。

因此，大部分状况下，yield都必须在try...except中，除非设计之初就是让yield以后的代码在with代码块内部出现异常时不执行。

测试代码：

以类的方式实现上下文管理器，当没有引起异常时， # 其执行结果与@contextmanager装饰器装饰器的上下文管理器函数相同，输出 [1, 2, 3, 4, 5]

    items_3 = [1, 2, 3]
    with ListTransaction(items_3) as working_3:
        working_3.append(4)
        working_3.append(5)
    print(items_3)

当执行代码过程当中引起异常时，即便没有对异常进行任何处理，__exit__方法也会正常执行，对self.orig_list进行修改（python是引用传值，而list是可变类型，对orig_list的任何引用的修改，都会改变orig_list的值），因此输出结果与没有引起异常时相同：[1, 2, 3, 4, 5]

    items_4 = [1, 2, 3]
    try:
        with ListTransaction(items_4) as working_4:
            working_4.append(4)
            working_4.append(5)
            raise RuntimeError('oops')
    except Exception as ex:
        print(ex)
    finally:
        print(items_4)

 

完整代码：https://github.com/blackmatrix7/python-learning/blob/master/class_/contextlib_.py