Python 模組介紹 - contextlib

在 Python 淺談 with 語句 一文中，我們已介紹基本的 with 語法與何謂 context manager 以及簡單的 context manager 實作。

除了該文所提供的實作方法之外，也可以利用 Python 內建的 contextlib 模組進行開發，此外該模組也提供若干實用的 context manager 可以使用。

本文環境 #

• Python 3.6.5

contextlib #

contextlib 提供以 decorator 的方式讓開發者能夠輕鬆製作 context manager:

from contextlib import contextmanager


@contextmanager
def simple_context():
    print('enter')
    yield
    print('exit')


with simple_context():
    print('I am in the context')

上述範例執行結果如下：

enter
I am in the context
exit

context manager 很適合應用在需要對某段程式做處置的情況，例如衡量某段程式的執行時間，就能夠利用類似以下的 context manager 進行衡量：

from contextlib import contextmanager
from time import perf_counter
from random import randint


@contextmanager
def profiling():
    s = perf_counter()
    yield
    print('time: ', perf_counter() - s)


with profiling():
    for i in range(1000):
        randint(0, 10)

上述範例執行結果如下：

time:  0.0013769530000331542

不過上述等利用 yield 而做成的 context manager 並無法被重複使用，例如：

from contextlib import contextmanager
from time import perf_counter
from random import randint


@contextmanager
def profiling():
    s = perf_counter()
    yield
    print('time: ', perf_counter() - s)


profiler = profiling()
with profiler:
    for i in range(1000):
        randint(0, 10)


with profiler:
    for i in range(1000):
        randint(0, 10)

上述範例執行第 2 次(也就是第 19 行)就會出現以下錯誤：

time:  0.0012313259940128773
Traceback (most recent call last):
  File "....py", line 19, in <module>
    with profiler:
  File "/.../versions/3.6.5/lib/python3.6/contextlib.py", line 83, in __enter__
    raise RuntimeError("generator didn't yield") from None
RuntimeError: generator didn't yield

原因在於該 context manager 其實是 1 個 generator, 一旦該 generator 一結束就無法再被使用，也因此該種 context managet 被稱為 single use context manager .

如果要讓 context manager 能夠重複使用，就得以類別實作 __enter__ 與 __exit__ 方法，例如以 class 形式修正前述範例後，該 profiler 就能夠正常執行多次：

class profiling(object):
    def __enter__(self):
        self.s = perf_counter()

    def __exit__(self, *exc):
        print('time: ', perf_counter() - self.s)


profiler = profiling()
with profiler:
    for i in range(1000):
        randint(0, 10)


with profiler:
    for i in range(1000):
        randint(0, 10)

上述執行結果：

time:  0.0013683729994227178
time:  0.0011790469980041962

其他內建的 context manager #

closing #

如果想讓某些有實作 close() 方法的類別，自動在離開 context manager 執行 close() 方法的話，可以使用contextlib.closing ，例如以下的 StringIO 在離開 closing context manager 時，自動執行了 close() 方法，因此導致第 8 行的 output.write('def') 無法正常寫入，因而發生錯誤：

import io
from contextlib import closing


with closing(io.StringIO()) as output:
    output.write('abc')

output.write('def')

Traceback (most recent call last):
  File "...", line 8, in <module>
    output.write('def')
ValueError: I/O operation on closed file

總的來說， contextlib.closing 相當適合需要確保 close() 方法一定得被呼叫的情境。

suppress #

有些時候我們可能希望在程式內忽略某些無關緊要的 exception, 例如以下是忽略 ZeroDivisionError 的寫法：

try:
    x = 1 / 0
except ZeroDivisionError:
    pass

同樣的寫法可以換成使用 contextlib.suppress ，程式整體看起來也相較簡潔:

from contextlib import suppress


with suppress(ZeroDivisionError):
    x = 1 / 0

上述範例會忽略 ZeroDivisionError 例外錯誤。

suppress 也可以忽略多種例外錯誤，例如：

from contextlib import suppress


with suppress(ZeroDivisionError, ValueError):
    x = 1 / 0

ExitStack #

contextlib 中有個有趣的 context manager 稱為 ExitStack .

該 context manager 的功用為讓開發者以堆疊(stack)的方式任意增加離開 context manager 時需要執行的函式(或者稱為 callback)。

也由於是 stack, 所以越晚增加的函示會先執行，例如以下範例：

from contextlib import ExitStack


def callback(*args, **kwargs):
    print('Callback', args[0])


def do_something():
    return True


with ExitStack() as stack:
    stack.callback(callback, 'a')
    stack.callback(callback, 'b')
    result = do_something()

上述範例執行結果如下，可以看到 stack.callback(callback, 'b') 最後被加入，也最先被執行：

Callback b
Callback a

另外， ExitStack 也可以選擇不執行 stack 內所有的 callback, 只要在離開 context manager 之前呼叫 pop_all() 即可，這也提供某種程度上的彈性，讓我們能夠控制程式流程：

from contextlib import ExitStack


def callback(*args, **kwargs):
    print('Callback', args[0])


def do_something():
    return True


with ExitStack() as stack:
    stack.callback(callback, 'a')
    stack.callback(callback, 'b')
    result = do_something()
    if result:
        stack.pop_all()

上述範例由於離開 context manager 之前呼叫 stack.pop_all() , 因此任何 callback 都不會執行。

以上就是關於 contextlib 的內容，官方文件中也紀錄其他本文未提及的 context manager, 可以花些時間翻閱，也許有些能夠運用在日常開發之中。

Happy coding!

References #

https://docs.python.org/3/library/contextlib.html