Pythonで引数の型に応じて返り値の型を変えるには@overloadを使う

## 環境
- Python 3.9.2
- VS Code 1.54.1
- [Pylance](https://marketplace.visualstudio.com/items?itemName=ms-python.vscode-pylance) 2021.3.1

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

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

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

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

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

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

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

のように使います。

```python
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")
```

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

https://docs.python.org/ja/3/library/typing.html#typing.overload

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

```python
# エラー
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"

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

```python
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()`](https://docs.python.org/ja/3/library/functions.html#open)の例を見てみます。
Pyrightの型定義では、modeの引数の値が`r`(テキストモード)か`rb`(バイナリモード)かに応じて型が変わります。

```python
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"

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/50391/ef0c6234-9eaa-de90-c623-3b7e70168f9d.png)

これは[`open()`の型定義が`@overload`で`mode`の引数の型ごとに定義されていて](https://github.com/microsoft/pyright/blob/1.1.120/packages/pyright-internal/typeshed-fallback/stdlib/builtins.pyi#L1067-L1152)、そのとおりに型推論されているからです。以下、Pyrightの型定義からの引用です

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

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

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

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

以下[wikipediaより引用](https://ja.wikipedia.org/wiki/%E5%A4%9A%E9%87%8D%E5%AE%9A%E7%BE%A9)

```java
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](https://www.typescriptlang.org/docs/handbook/2/narrowing.html)を参照）

```typescript
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で挙がっているようです](https://www.python.org/dev/peps/pep-0647/)が、すでにPyrightには実装されているようです？

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

## 参考
- [typing \-\-\- 型ヒントのサポート — Python 3\.9\.2 ドキュメント](https://docs.python.org/ja/3/library/typing.html)
- [Pylance \- Visual Studio Marketplace](https://marketplace.visualstudio.com/items?itemName=ms-python.vscode-pylance)
- [microsoft/pyright: Static type checker for Python](https://github.com/microsoft/pyright)