【Python】関数の引数の型に応じて戻り値の型を動的に変える型ヒントの書き方

Pythonでもジェネリクスみたいに型を書きつつそれを動的に扱えないかな～と思ってガチャガチャやって辿り着いたものを共有

環境

• Python 3.11.9

Python 3.12以降でより良い書き方あるかは不明

例

from typing import TypeVar
​
T = TypeVar("T", int, int | None)
​
def init_list(initial_value: T, num: int) -> list[T]:
    # initial_value: T = 0 のようにデフォルト値も入れられる
    res: list[T] = [initial_value] * num
    return res

init_list(0, 1)  # 戻り値の型は`list[int]`になる
init_list(None, 1)  # 戻り値の型は`list[int | None]`になる

そもそもTypeVarとは何か、から説明します。

簡単に言うと「Anyに少し制約を加えられ、実際の値の型になってくれるもの」という感じです。

つまり、Anyはどんな値が渡ってこようともAny型でしかありませんが、
TypeVarは例えばint型の値が渡ってきたらint型になってくれます。

よくUnionと一緒に名前が出されますが、

• Unionは、AもBも同時に許容する
• TypeVarはAかBを許容する

といった違いがあります。

実際に使ってみた例を以下に示します。

# 両方の型を混ぜて使える
x: list[int | None] = [None, 0, 1]
​
U = TypeVar("U", int, None)
​
def to_list(*values: U) -> list[U]:
    return list(values)
​
y1 = to_list(None, 0, 1) # できない
y2 = to_list(None, None) # できる
y3 = to_list(0, 1) # できる

なんならy3 = to_list(0, 1)のときy3の型はintに絞られます。
(この場合は後続の処理で「y3にNoneが含まれていたら～」といった判定も不要になる)

ちなみに最初に挙げた例では、このTypeVarとUnionを組み合わせを実現しています。

意外とうまいことやってくれるみたい