Pythonのcontext managerとgeneratorについて

背景

時たま必要になるpythonの細かい所メモ

context manager

context managerの目的

__enter__ と __exit__ を持つオブジェクトがコンテキストマネージャ
with文で「入る・出る」を強制できる
ファイルclose、接続切断、ロック解放、トランザクション終了などに使う
goでいうdeferに近い

1
2
3
with open("a.txt") as f:
    data = f.read()
# ここに来た時点で自動で f.close() される

同期context managerの例

I/F的には、enterとexitのメソッドを持つクラス。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class MyContext:
    def __enter__(self):
        print("入室")
        return self # これがwith MyContext() as xのxになる

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("退室")

# 使い方
with MyContext():
    print("作業中")

# output
# 入室
# 作業中
# 退室

同期context managerを関数で作る

関数でもcontext managerは作れる。

1
2
3
4
5
6
7
8
9
from contextlib import contextmanager

@contextmanager
def my_context():
    print("入る")
    try:
        yield
    finally:
        print("出る")

generatorはデコレーターでcontext managerにもできる。

※:

with ... as x: の xはyieldの値
yield は必ず1回
yield 前 = enter
yield 後 = exit
yield しない / 複数回 → エラー

非同期のcontext manager

__aenter__、__aexit__を実装する
async withで使う

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
import asyncio

class AsyncResource:
    async def __aenter__(self):
        print("接続開始")
        await asyncio.sleep(1)  # 非同期I/O
        return self

    async def __aexit__(self, exc_type, exc, tb):
        print("接続終了")
        await asyncio.sleep(1)

使用:

1
2
async with AsyncResource() as r:
    print("使用中")

同期と非同期のキーワードの違い

種類	構文	中で使うメソッド
同期	`with`	`__enter__`, `__exit__`
非同期	`async with`	`__aenter__`, `__aexit__`

exitstack

終了処理の順番を担保したもの。

1
2
3
4
5
6
7
from contextlib import ExitStack

with ExitStack() as stack:
    if use_http:
        session = stack.enter_context(make_session())   # 片付けを自動登録
    if use_lock:
        lock = stack.enter_context(acquire_lock())      # 片付けを自動登録

こうすることで、lockを外して、sessionを閉じるを順序よくできる。つまり、次のコードを自動化したようなもの。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
session = None
lock = None

try:
    if use_http:
        session = make_session()
    if use_lock:
        lock = acquire_lock()
    # 何か処理…
finally:
    if lock:
        lock.release()
    if session:
        session.close()

generator

generatorの目的

「値を一気に作らず、必要になったら順に作って返す仕組み」（遅延評価）
yield を使う関数、またはジェネレーター式で作る
1個ずつ順番に値を返せる（巨大データや無限列に強い）

同期generator

1
2
3
4
5
6
7
8
9
def count_up():
    n = 0
    while True:
        yield n
        n += 1

g = count_up()
next(g)  # 0
next(g)  # 1

非同期generator

定義

1
2
3
4
async def agen():
    for i in range(3):
        await asyncio.sleep(1)
        yield i

これを回すにはasync forを使う。

1
2
async for x in agen():
    print(x)

同期と非同期のキーワードの違い

種類	生成	ループ
同期	`yield`	`for`
非同期	`yield` + `await`	`async for`

同期generatorをクラスで作る

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class CountUp:
    def __init__(self, limit):
        self.n = 0
        self.limit = limit

    def __iter__(self):
        return self

    def __next__(self):
        if self.n >= self.limit:
            raise StopIteration
        value = self.n
        self.n += 1
        return value

使い方:

1
2
3
for x in CountUp(3):
    print(x)
# 0, 1, 2

__iter__ → イテラブルを返す
__next__ → 1ステップ進める
StopIteration → 終了シグナル