環境

Python 3.9.2
VS Code 1.54.1
Pylance 2021.3.1

概要

下記のように、引数の型に応じて返り値の型も変わる関数を定義します。

def double(arg: Union[int, str]):
    return arg * 2

このままだと、この関数の返り値の型は何を渡した時もUnion[int, str]と推論されてしまいます。
そのため、返り値をintやstr型の変数に代入すると型チェックでエラーが出てしまいます。

# エラーが出る
var2: int = double(1)
var11: str = double("1")

このように引数の型に応じて返り値の型を変えたい時は、@overloadを使って引数の型ごとに返り値を定義してあげることでエラーを解消できます。

`@overload`を使う

Pythonにおける@overloadは、型チェッカーのためだけの関数定義です。付いた関数を実行時に呼び出すことはできません。

@overloadを付けて、型チェッカーのために型定義だけの関数を記述する
@overloadを付けずに、実際に実行される関数を記述する

のように使います。

from typing import overload

@overload
def double(arg: int) -> int:
    ...
@overload
def double(arg: str) -> str:
    ...

def double(arg: Union[int, str]):
    return arg * 2

# エラー出ない
var2: int = double(1)
var11: str = double("1")

詳細は公式ドキュメントを参照してください。

@overloadはちゃんと関数内部の返り値の型チェックもしてくれるようです。

# エラー
def double(arg: Union[int, str]):
    return str(arg) * 2

このように文字列型しか返さない定義にしてしまうと、ちゃんとint型と互換性が無いというエラーになります。

Overloaded function implementation is not consistent with signature of overload 1
Function return type "int" is incompatible with type "str"
"int" is incompatible with "str"

別解: ジェネリクス

上記のような、引数と同じ型を返すぐらいの例であれば、ジェネリクスを使うことでも型チェックのエラーを回避できます。

from typing import TypeVar

T = TypeVar('T', int, str)

def double(arg: T) -> T:
    return arg * 2

# エラー出ない
var2: int = double(1)
var11: str = double("1")

`open()`関数の例

実際に、Pythonの組み込み関数であるopen()の例を見てみます。
Pyrightの型定義では、modeの引数の値がr(テキストモード)かrb(バイナリモード)かに応じて型が変わります。

import pickle
# エラーが出る
with open('hoge.pickle', mode='r') as f:
    a = pickle.load(f)

# 正しいコード
with open('hoge.pickle', mode='rb') as f:
    a = pickle.load(f)

上記のようなコードを書くと、下記のようなエラーが表示されます。

Argument of type "TextIOWrapper" cannot be assigned to parameter "file" of type "IO[bytes]" in function "load"
TypeVar "AnyStr@IO" is invariant
"str" is incompatible with "bytes"

これはopen()の型定義が@overloadでmodeの引数の型ごとに定義されていて、そのとおりに型推論されているからです。以下、Pyrightの型定義からの引用です

builtins.pyi

@overload
def open(
    file: _OpenFile,
    mode: OpenTextMode = ...,
# ---- 中略 ----
) -> TextIOWrapper: ...

builtins.pyi

@overload
def open(
    file: _OpenFile,
    mode: OpenBinaryModeReading,
# ---- 中略 ----
) -> BufferedReader: ...

それぞれOpenTextModeは"r"、OpenBinaryModeReadingは"rb"などを含むリテラル型です。

overloadについて所感

overload（多重定義）というと、Java等では単に「同名のメソッドを複数定義する仕組み」として使われてきました。

以下wikipediaより引用

public class Main {
    static void testMethod(String str) { System.out.println("String version is called."); }
    static void testMethod(Object obj) { System.out.println("Object version is called."); }
    public static void main(String[] args) {
        String str = "test";
        Object obj = str;
        testMethod(str); // String バージョンが呼ばれる。
        testMethod(obj); // Object バージョンが呼ばれる。
    }
}

一方PHP等では、引数の型などによって分岐することによって実質的にoverloadと同じインターフェースを実現してきました。TypeScriptの型推論はこの分岐による型の絞り込みを読み取ることができました。（type narrowingを参照）

function double(arg: number | string) {
    if (typeof arg === "number") {
        // ここではargの型はnumberだけと推論される
        return arg * 2;
    }
    // ここではargの型はstringだけと推論される
    return arg + arg;
}

Pythonは、overloadを注釈のようなランタイムに使われない形で記述する所はTypeScriptらしいと思いました。メソッドの定義上はunion型となっていてもちゃんとoverloadの定義に沿って型チェックでエラーを出してくれる点も使いやすいです。またTypeScriptのようなtype narrowingについては2021年3月13日時点でPython 3.10のProposalがDraftで挙がっているようですが、すでにPyrightには実装されているようです？

def double(arg: Union[int, str]):
    if isinstance(arg, int):
        # ここではint型と推論される
        return arg * 2
    # ここではstr型と推論される
    return arg + arg

環境

概要

@overloadを使う

別解: ジェネリクス

open()関数の例

overloadについて所感

参考

`@overload`を使う

`open()`関数の例