使用@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