Python
typing

【Python】戻り値の型がわからないメソッドを正しく呼びたい【typing】

冷静に考えれば当たり前の話だけど、ちょっと迷ったのでメモ。

要件

foo()メソッドを持つオブジェクトを受け取り、そのオブジェクトのfoo()メソッドを呼んで、その戻り値を返す関数call_foo()を定義したい。

fooables.py
class Fooable1:
    def foo(self) -> int:
        return 42

class Fooable2:
    def foo(self) -> str:
        return 'bar'
call_foo.py
def call_foo(x: '???') -> '???':
    return x.foo()
main.py
from fooables import Fooable1, Fooable2
from call_foo import call_foo

if __name__ == '__main__':
    a = call_foo(Fooable1())  # Expected type for a: int
    b = call_foo(Fooable2())  # Expected type for b: str
    c = call_foo(42)  # Expected to raise error at static type checking

call_foo()の引数xfoo()を持つことだけが規定されていて、x.foo()の戻り値の型をcall_foo.pyモジュールから見ることはできない。

失敗例

call_foo.py
from typing import Any

def call_foo(x: Any) -> Any:
    return x.foo()

Anyを受け取って、Anyを返すようにしてみた。

main.py
from fooables import Fooable1, Fooable2
from call_foo import call_foo

if __name__ == '__main__':
    a = call_foo(Fooable1())  # Inferred type is "Any"
    b = call_foo(Fooable2())  # Inferred type is "Any"
    c = call_foo(42)  # No error and inferred as "Any"

全てAnyになってしまい、何の情報も残らない。失敗。

正解例

call_foo.py
from typing_extensions import Protocol
from typing import Generic, TypeVar
from abc import abstractmethod

T_co = TypeVar('T_co', covariant=True)

class FooableType(Generic[T_co], Protocol):
    @abstractmethod
    def foo(self) -> T_co:
        raise NotImplementedError()

def call_foo(x: FooableType[T_co]) -> T_co:
    return x.foo()

これで期待通りの動作になる。

main.py
from fooables import Fooable1, Fooable2
from call_foo import call_foo

if __name__ == '__main__':
    a = call_foo(Fooable1())  # Inferred type is "int"
    b = call_foo(Fooable2())  # Inferred type is "str"
    c = call_foo(42)  # Error: Argument 1 to "call_foo" has incompatible type "int"; expected "FooableType"

解説

foo()を持つ任意の型」を表現するためには、Protocolを使うのが良い。
foo()を持つクラスを作ってそれを継承してもらってもいいが、今回の場合Fooableがどのようなクラスなのかcall_fooモジュールは知らないので、こちらの要求するクラスを継承してもらえるとは考えにくい。
従って、Protocolabstractmethodを用いて表現することになる。

foo()の返す型」を返したい場合は、ジェネリッククラスGeneric[]と型変数TypeVarを用いて、 C++ で言うところのテンプレートクラスにしてしまうのが良い。
Generic[T]を継承することで、このクラスをジェネリッククラス、すなわち受け取った型によって引数や戻り値などの型を変更できるクラスにすることができる。
この型は今回戻り値の型にしか使わないので、共変であること(例えば、FooableType[float]型の変数にFooableType[int]型の値を代入できること)を明示すると良い。実際、mypy 0.630 はこのT_coを不変にすると「Invariant type variable 'T_co' used in protocol where covariant one is expected」と怒ってくれる。

おわり

ProtocolGenericの組み合わせは強力で、かなり任意のお気持ちを表すことができる。
積極的に使っていきたい。