Pythonで@overloadした関数を呼んでリストに入れると型チェックでエラー

Question

下記コードでエラーが出る原因を知りたいです。

下記の中で、hoge は、引数 type_ に応じて str か int を返します。
fuga の中でこれを呼び出してリストの中に入れると、型チェックでエラーになります。

@overload
def hoge(number: int, type_: Literal["STRING"]) -> str:
    ...

@overload
def hoge(number: int, type_: Literal["INTEGER"]) -> int:
    ...

def hoge(number: int, type_: Literal["STRING", "INTEGER"]) -> str | int:
    if type_ == "STRING":
        return str(number)
    else:
        return number

def fuga(type_: Literal["STRING", "INTEGER"]):
    piyo: list[str] | list[int] = [hoge(i, type_) for i in range(3)]  # ここでエラー
    print(piyo)

エラーの内容

Expression of type "list[str | int]" cannot be assigned to declared type "list[str] | list[int]"

実際に fuga() を実行すると、(当然ながら)下記の通りです。

>>> fuga("STRING")
['0', '1', '2']
>>> fuga("INTEGER")
[0, 1, 2]

fuga の中で type_ は "STRING" か "INTEGER" のどちらか一方に決まっているので、 piyo の型は list[str | int] ではなく list[str] か list[int] のどちらか一方のはずです。

このエラーを解決する方法はありますでしょうか？

（list[str | int] は、例えば [0, '1', 2] のようにリストの要素に str と int が混ざったものを含むと理解していますが合っていますでしょうか？）

Answer 1

fuga の中で type_ は "STRING" か "INTEGER" のどちらか一方に決まっているので、 piyo の型は list[str | int] ではなく list[str] か list[int] のどちらか一方のはずです。

それはPython処理系で検知できない、ということかと思います。

このエラーを解決する方法はありますでしょうか？

castでlist[str] | list[int]型を与えてください。

Comments

[naohash]

ありがとうございます。
castするしかないということですが、もし、下記よりスマートな書き方があれば教えていただけるとありがたいです。
無駄に行数を使って読みづらくなるなぁと。

piyo = [hoge(i, type_) for i in range(3)]
if type_ is Literal["STRING"]:
piyo = list(map(str, piyo))
else:
piyo = list(map(int, piyo))

markdownがコメントで使えず読みづらくて申し訳ありません。

[melian]

> naohash さん

piyo: list[str] | list[int] = cast(list, [hoge(i, type_) for i in range(3)])

としてみてはどうでしょうか。直後に reveal_type(piyo) で確認すると、

note: Revealed type is "Union[builtins.list[builtins.str], builtins.list[builtins.int]]"

となっています。

[naohash]

遅くなって申し訳ありません。
これは知りませんでした。ありがとうございます！