Python: コンテキストマネージャの型ヒント

概要

with構文で使われる、任意のコンテキストマネージャはcontextlib.AbstractContextManager型の変数に代入できる。​AbstractConrextManagerを継承している必要はない​。(mypy v0.942で確認)

AbstractConrextManagerを継承していないのであればエラーになるべきで、次のようなプロトコル型が必要になると思うのだがどうなんだろう？

T = TypeVar("T", covariant=True)

class ContextManager(Protocol[T]):
    def __enter__(self) -> T:
        ...

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        ...

具体例

contextlib.nullcontextの使用例にある次の関数に型ヒントを付けたい。

from contextlib import nullcontext

def process_file(file_or_path):
    if isinstance(file_or_path, str):
        # file_or_pathが文字列ならファイルを開く
        cm = open(file_or_path)
    else:
        # ファイルオブジェクトなら呼び出し側が閉じるべき
        cm = nullcontext(file_or_path)

    with cm as file:
        # ファイルから読み込み

openがテキストモードでファイルを開いているので、file_or_pathはstr | TextIOだと仮定する。この時cmの型ヒントはどうすべきか？

なお、何も型ヒントを与えない場合、

error: Incompatible types in assignment (expression has type "nullcontext[TextIO]", variable has type "TextIOWrapper")

というエラーが出る。

解決策

typingパッケージを見ると、typing.ContextManagerという探していた型が定義されているが、v3.9以降はcontextlib.AbstractContextManagerを使えと書いてある。

AbstractContextManagerを使って次のようにアノテーションするとエラーは解消された。

from contextlib import AbstractContextManager, nullcontext

def process_file(file_or_path: str | TextIO) -> None:
    cm: AbstractContextManager[TextIO]
    if isinstance(file_or_path, str):
        # file_or_pathが文字列ならファイルを開く
        cm = open(file_or_path)
    else:
        # ファイルオブジェクトなら呼び出し側が閉じるべき
        cm = nullcontext(file_or_path)

    with cm as file:
        # ファイルから読み込み

疑問

コンテキストマネージャの型にAbstractContextManagerを使うと、AbstractContextManagerを継承していない野良コンテキストマネージャ( ​__enter__と__exit__を実装しているだけのクラス)は代入できないのでは？と思い、実験をしてみた。

# AbstractContextManagerを継承していないコンテキストマネージャ
class MyContextManager:
    def __enter__(self) -> TextIO:
        ...

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        ...

cm: AbstractContextManager[TextIO]
cm = MyContextManager()

予想に反してエラーはでませんでした。AbstractContextManagerには継承関係にないコンテキストマネージャもなぜか代入できるようです。

そもそもTextIOがAbstractContextManagerを継承していないようなので、この辺りの都合でエラーにならないのかも知れません。

typing.py

class TextIO(IO[str]):
    ...

class IO(Generic[AnyStr]):
    ...

class Generic:
    ...

本来であれば、次のようなプロトコルクラスが必要だと思うので、将来挙動が変更された時のために記しておきます。

from types import TracebackType
from typing import Optional, Protocol, TypeVar

T = TypeVar("T", covariant=True)

class ContextManager(Protocol[T]):
    def __enter__(self) -> T:
        ...

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        ...