Context Manager是一種可以讓我們使用with，於進出某段程式碼時，執行某些程式碼的功能。

Context Manager Protocol

Context Manager Protocol要求需實作__enter__及__exit__兩個dunder method。

__enter__

__enter__的signature如下：

__enter__()

__enter__不接受參數，其返回值將可以用with搭配as的語法取得，例如with ctxmgr() as obj。

__exit__

__exit__的signature如下：

__exit__(exc_type, exc_val, exc_tb)

其接收三個參數：

• exc_type為例外的class。
• exc_val為例外的obj（或想成exc_type的instance）。
• exc_tb為一個traceback obj。

當__exit__回傳值為：

• truthy時（bool(回傳值)為True），會忽略例外。
• falsey時（bool(回傳值)為False），會正常報錯。由於當function沒有顯性設定回傳值時，會回傳None。而None是falsey，所以context manager預設情況為正常報錯。

基本型態

Context Manager一般有兩種型態：

• 型態1是希望在進入時啟動資源，而在離開時關閉資源。常見的應用場景是開關檔案，建立database client、ssh client或http client等等。
• 型態2是希望能在context manager下，「暫時」有些特別的行為。常見的應用場景是設定臨時的環境變數或是臨時的sys.stdout或sys.stderr。

型態1

型態1接收的參數，通常用來生成底層真正使用的obj。例如建立一個PostgreSQL的connection可能需要host、port、database name、username及password等等參數。

於__enter__中可以做一些setup，例如建立connection、進行logging等。至於返回值一般會返回self，因為這樣可以方便使用於class中的其它function，但依照使用情況的不同，有時候返回底層obj會更加方便。

於__exit__中可以做一些cleanup，例如關閉connection、進行logging等。此外，有可能需要處理遇到的例外，並決定返回truthy或falsey。

# 01 PSEUDO CODE!!!
class Object:
    def __init__(self, **kwargs): ...
    def start(self): ...
    def finish(self): ...


class MyContextManager:
    def __init__(self, **kwargs):
        self._kwargs = kwargs

    def _make_obj(self, **kwargs):
        return Object(**kwargs)

    def setup(self):
        """set up something and possibly call self._obj.start() to do something"""
        self._obj = self._make_obj(**self._kwargs)
        self._obj.start()

    def cleanup(self):
        """Possibly call self._obj.finish() to do something and clean-up something"""
        self._obj.finish()

    def __enter__(self):
        self.setup()
        # can:
        # 1. return self
        # 2. return self._obj
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # may need to handle exceptions
        self.cleanup()

型態2

型態2通常只接收單個或少數參數，這些參數可以用來建構於context manager中「暫時」想要的行為。例如redirect stdout，或是暫時覆寫某些環境變數等。

於__enter__中，我們會先使用getter儲存當前的狀態，再使用setter實現想要的行為。至於返回值，要看當前應用的情況，即使不返回（即返回None）也是常見的情況。

於__exit__中，我們再使用setter回復原先的狀態。一樣需視情況來處理遇到的例外，並決定返回truthy或falsey。

# 02 PSEUDO CODE!!!
class MyContextManager:
    def __init__(self, new_x):
        self._new_x = new_x
        self._x = 'x'

    def __enter__(self):
        self._old_x = self._x
        self._x = self._new_x
        # can:
        # 1. return self
        # 2. return self._new_x
        # 3. return None (implicitly)
        return self._new_x

    def __exit__(self, exc_type, exc_val, exc_tb):
        # may need to handle exceptions
        self._x = self._old_x  # back to original state
        del self._old_x  # delete unused variable

三種context manager類型

Context Manager可以分為single use、reusable及reentrant三種類型。

single use

single use是最常用的類型。每次需要使用這類型的context manager都需重新建立，重複使用將會raise RuntimeError。建議的使用方法是，盡量使用with MyContextManager as ctx_mgr的語法，而不要將其先存在一個變數，例如ctx_mgr = MyContextManager()，然後再with ctx_mgr，來降低發生重複使用的機率。

reentrant

reentrant是指在with ctx區塊內再產生一個以上的with ctx區塊。redirect_stdout與redirect_stderr即是此種類型，我們稍後會欣賞其源碼。

reusable

reusable是排除有reentrant特性的context manager。其可以多次呼叫，但是如果將其當reentrant來使用時，會報錯或出現非預期的行為。

ContextDecorator and contextmanager

ContextDecorator

假如您有一個實作了__enter__及__exit__的context manager，那麼只要再繼承ContextDecorator，這個context manager就能當作decorator使用。其源碼非常精簡，就像是附加一個__call__在context manager上。其功能是在被裝飾的function被呼叫時，會自動將該function包在with區塊內執行，就像是顯示使用with一樣，真是一個巧妙的設計呀。

class ContextDecorator(object):

    def _recreate_cm(self):
        return self

    def __call__(self, func):
        @wraps(func)
        def inner(*args, **kwds):
            with self._recreate_cm():
                return func(*args, **kwds)
        return inner

contextmanager

當使用contextmanager裝飾在一個generator function上時，此generator function將具有context manager的特性，且其也可以作為decorator使用（因為contextmanager內部實作有使用ContextDecorator）。

下面是Python文件的示例。

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)

其中yield的resource就相當於是__enter__中回傳值，可以方便我們使用下方with managed_resource as 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

redirect_stdout與redirect_stderr源碼

contextlib內有不少實作了context manager的好用工具，我們一起來瞧瞧redirect_stdout與redirect_stderr是怎麼實作的。

class redirect_stdout(_RedirectStream):
    _stream = "stdout"
    
class redirect_stderr(_RedirectStream):
    _stream = "stderr"

原來兩個都是繼承_RedirectStream而來，只是_stream這個class variable設的不同而已，讓我們再繼續追下去。

class _RedirectStream(AbstractContextManager):

    _stream = None

    def __init__(self, new_target):
        self._new_target = new_target
        # We use a list of old targets to make this CM re-entrant
        self._old_targets = []

    def __enter__(self):
        self._old_targets.append(getattr(sys, self._stream))
        setattr(sys, self._stream, self._new_target)
        return self._new_target

    def __exit__(self, exctype, excinst, exctb):
        setattr(sys, self._stream, self._old_targets.pop())

_RedirectStream於：

• __init__中，接收一個參數，為想要redirect的新目標。另外建立了一個self._old_targets的list來收集舊目標。
• __enter__中，將當前的sys.stdout或sys.stderr附加到self._old_targets後，返回self._new_target（不是self）。這麼一來，我們就可以在as的關鍵字後，得回self._new_target。
• __enter__中，將當前的sys.stdout或sys.stderr設為self._old_targets所pop出來的值。list的pop可以同時刪除最後一個元素並將其返回，用在此處可謂恰如其分。

_RedirectStream屬於我們的型態2，於__enter__中儲存當前狀態後，改變到新狀態，最後再於__exit__中恢復原來狀態。而且其註解也寫明其是re-entrant的，這也是為什麼我們需要self._old_targets幫忙來儲存一個以上的狀態。

參考資料

• Python Morsels - Creating a context manager in Python。
• Python for network engineers - VI. Basics of object-oriented programming - 23. Special methods - Protocols - Context Manager。

Code

本日程式碼傳送門。