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