Zenn
🐍

Pythonのコンテキストマネジャーをまとめる

に公開
Python
tech

始めに

Pythonのコンテキストマネージャは、with文と組み合わせてリソース（ファイル、DB接続、ロックなど）の確保と解放を自動化する仕組みです。__enter__メソッドでセットアップ、__exit__メソッドでクリーンアップを行うことで、例外発生時も安全にリソース管理できます。

知っておけば便利なことが多いので、ブログに便利な箇所をまとめます。

環境

  • Python
    • 3.13.3

実装

基本的には明示的にリソースの開放が必要な時に使用します。ただ、DB自体のリソース開放はライブラリ側が自動的に行ってくれますので、インスタンス生成時にコンテキストマネジャーを使用しましょう。

async def get_reader_db():
    async with async_session() as session:
        yield session
        
SessionReaderDep = Annotated[AsyncSession, Depends(get_reader_db)]

@router.get("/sample")
async def sample(
    db: SessionReaderDep,
) -> Any:
    try:
        // 適当な処理
        await sample2(db);
        await session.commit()
    except HTTPException:
        await session.rollback()
        raise
    except Exception as e:
        await session.rollback()
        raise e

ただし、通常時のコミット処理や、エラーが発生したときのロールバックやログの出力等々はアプリケーション側の責任となります。どの箇所でも同じ記載をすることになりますので、それもまたコンテキストマネジャーで共通化しましょう。

from contextlib import asynccontextmanager

from fastapi import HTTPException
from sqlalchemy.ext.asyncio import AsyncSession

# コンテキストマネジャーの定義
@asynccontextmanager
async def session_context(session: AsyncSession) -> None:
    try:
        yield session
        await session.commit()
    except HTTPException:
        await session.rollback()
        raise
    except Exception as e:
        await session.rollback()
        raise e
 
 
@router.get("/sample")
async def sample(
    param_db: SessionReaderDep,
) -> Any:
    async with session_context(param_db) as db:
        // 適当な処理
        await sample2(db);

さらに重要なのが__exit__または__aexit__処理でFalseかNoneを返却することです。もし定義しない場合、エラー発生行がコンテキストマネジャーで発生したことになり、何行目でエラーが発生したかがわからなくなります。ライブラリとしては意図的にエラーをまとめていると思いますが、アプリケーションの場合にはエラー行を特定させない行為にメリットは薄いので、必ず追加で定義したほうが良いです。

例外は、別の例外を送出するような finally 節が無い場合にのみ呼び出しスタックへ伝わります。新しい例外によって、古い例外は失われます。
`

@asynccontextmanager
async def session_context(session: AsyncSession) -> None:
    try:
        yield session
        await session.commit()
    except HTTPException:
        await session.rollback()
        raise
    except Exception as e:
        await session.rollback()
        raise e
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):  # type: ignore
        """
        呼出元にエラーが発生したときに伝播させたいので、Falseを常に返却する。
        発生しないと、全てのエラーがcontext_managerの行で発生したことになる。
        """
        return False

@router.get("/sample")
async def sample(
    param_db: SessionReaderDep,
) -> Any:
    async with session_context(param_db) as db: // この行でエラーになったことになる
        // 適当な処理
        await sample2(db);
        
async def sample2(db):
    raise ValueError("") // 本当はこの行でエラーになったのに。

また、コンテキストマネジャーはデコレータとしても使用できます。例えば、バッチ処理で処理開始前と処理開始後にログを出力して、処理時間を計りたい等々があったとします。それも簡単に作成できます。

@asynccontextmanager
async def timeit(operation_name: str = "Operation"):
    """
    処理時間を計測し、開始と終了時にログを出力するコンテキストマネージャー

    Args:
        operation_name: 操作名（ログに表示される）
    """
    start_time = time.time()
    logger.info(f"[{operation_name}] 処理開始")

    try:
        yield
    finally:
        end_time = time.time()
        elapsed_time = end_time - start_time
        logger.info(f"[{operation_name}] 処理完了 - 実行時間: {elapsed_time:.4f}秒")
        
async def test_decolator():
    @timeit("TEST DECOLATOR")
    async def test_func():
        await asyncio.sleep(0.1)

    # WHEN
    await test_func()
    # ログとして、次のログが出力される
    # [TEST DECOLATOR] 処理開始
    # [TEST DECOLATOR] 処理完了 - 実行時間: 0.1022秒

このようにリソース管理を安全にするだけでなく、簡単にデコレータ作成までできてしまうので、意外と便利な機能です。

ソースコード

終わりに

正直なところでいえば、コンテキストマネジャーに絞った記事というよりは、__exit__による挙動に少々苦しめられたのでそれを共有するための記事でした。ただ、内容が薄くなってしまったので、コンテキストマネジャーをまとめるという形で記事を記載しています。

参考情報

Discussion