Pythonの型アノテーションテクニック（1/N）

Pythonの型アノテーションについてのメモ。
ライブラリ実装や仕事で使うコードなど使い回す類のコードを書くときに、型アノテーションがあると色々捗ります。
適切に型アノテーションが成されていれば、テキストエディタやIDEはいい感じに補間してくれたり、引数リストを表示してくれます。
一方で型アノテーションにフォーカスした記事は少ないので簡単にまとめます。

思い出したら続きも書きます。

基本

色々なアノテーションの書き方の内、最もベーシックなものです。
コロンの後ろに型を続けて記述することを基本とします。
関数の戻り値だけは、-> (アロー)に続けて記述します。

# 変数の型アノテーション
name: str = "John"
age: int = 30
is_student: bool = False
height: float = 5.9

# 関数の型アノテーション
def hello(name: str) -> str:
    return "Hello, " + name

def add(x: int, y: int) -> int:
    return x + y

# 定数のアノテーション
from typing import Literal, Any
# Literal: 定数。
# Any: 型が不明な場合。
def test(mode: Literal["w", "r", "a", "w+", "r+"])->Any:
    return open("test.txt", mode)

また、listやdictは以下のように記述します。

from typing import Dict, List

# 辞書の型アノテーション
person: Dict[str, int] = {'age': 30, 'height': 175, 'weight': 70}
# キーは文字列で、値は整数です。

# リストの型アノテーション
numbers: List[int] = [1, 2, 3, 4, 5]
# リストの要素はすべて整数です。

# ネストされた辞書やリストの型アノテーション
data: Dict[str, List[int]] = {'A': [1, 2, 3], 'B': [4, 5, 6]}

ちなみに、Python 3.9以上では以下のように簡潔に記述することができます。
特に、import文も必要ありません。

# 辞書の型アノテーション
person: dict[str, int] = {'age': 30, 'height': 175, 'weight': 70}

# リストの型アノテーション
numbers: list[int] = [1, 2, 3, 4, 5]

# ネストされた辞書やリストの型アノテーション
data: dict[str, list[int]] = {'A': [1, 2, 3], 'B': [4, 5, 6]}

Python 3.8のEOL（End Of Life=サポート終了）は2024年10月に予定されているので、今から新規でプロジェクトを立てるのであれば、Python 3.9以上を想定して上記の記法を採用すれば良いと思います。
もし、Python 3.8に対応したい場合も以下のimportを記述すれば同様に利用可能です。

from __future__ import annotations

TYPE_CHECKING

最も基本的なテクニックです。
TYPE_CHEKINGは型チェッカーが読み出す時のみTrueとなり、普段実行する時はFalseとなる変数です。

例えば、型ヒントを記述するためだけに、importに時間がかかる激重モジュールをimportするとプログラムを実行するまでに要する時間が伸びてしまう。
現実、TensorFlowやPyTorchを使用するPythonのCUIプログラムを作ったときに、プログラムのヘルプが表示されるまでに尋常じゃない時間がかかる経験をした人はたくさん居ると思います。
これを回避する例が以下です。なお、heavy_moduleはimportするのに10秒程要するライブラリとします。

import heavy_module
def test() -> heavy_module.HeavyItem:
    return heavy_module.HeavyItem

import typing
if typing.TYPE_CHECKING:
    import heavy_module
    def test() -> heavy_module.HeavyItem:
        # 型チェック用の関数はEllipsis演算子（...）を配置することが一般的
        ...
else:
    def test():
        # 関数が呼ばれたタイミングでimport
        # 初めてtestが呼ばれた時だけ実行に10秒ほど時間がかかる
        import heavy_module 
        return heavy_module.HeavyItem()

その他にも、何か別のライブラリのラッパ関数を作る場合にも有効です。
例では引数が3つしかありませんが、引数が10も20もある関数をラップする場合に、ライブラリの関数の型アノテーションを見えるようにしたい場合に有効です。

import typing

def library_func(x: int, y: int, z: int) -> int:
    return x + y + z

if typing.TYPE_CHECKING:
    # ライブラリから型定義をコピーしてくる。型ヒントの表示はこの関数を参照する
    def test(x: int, y: int, z: int) -> int:
        ...
else:
    # ライブラリの更新に伴ってlibrary_funcの引数が増えてもすべての引数を渡すことができる。
    def test(*args, **kwargs):
        return library_func(*args, **kwargs)

overload

関数のオーバーロードです。
オーバーロードと言っても、一般的な言語のオーバーロードと違って、Pythonでは型ヒントのためだけに使用されます。
なお、overloadはTYPE_CHECKINGで囲うことが一般的です。

import typing
from typing import overload
if typing.TYPE_CHECKING:
    @overload
    def test(x: int) -> int:
        ...

    @overload
    def test(x: int, y: int) -> int:
        ...

def test(x, y=None):
    if y is None:
        return x
    else:
        return x + y

test(10)
test(10, 20)

正直この程度のユースケースであれば、overloadをわざわざ定義する必要がありません。

有効なユースケースは入力に応じて戻り値が変わるケースや、引数の与え方に組み合わせがある場合です。
具体例は以下の通りです。

例1

import typing
from typing import overload
if typing.TYPE_CHECKING:
    # 引数にintが与えられたら戻り値がintになる
    @overload
    def test(x: int) -> int:
        ...

    # 引数にfloatが与えられたら戻り値がfloatになる
    @overload
    def test(x: float) -> float:
        ...

def test(x: int | float) -> int | float:
    return x

例2

import typing
from typing import overload, Literal

if typing.TYPE_CHECKING:
    import torch
    import numpy as np

    # 与えられた引数の値によって必須とする引数を切り替えるようなアノテーションも記述可能
    # parallelに"gpu"が与えられたらどのgpuを使用するかを引数に与えてほしい
    @overload
    def test(x: torch.Tensor, parallel: Literal["gpu"], gpus: list[int]) -> int:
        ...

    # parallelに"cpu"が与えられたら使用するCPUの個数を引数に与えてほしい
    @overload
    def test(x: np.ndarray, parallel: Literal["cpu"], num_cpu: int) -> float:
        ...

# 実際の実装
def test(
    x: torch.Tensor | np.ndarray,
    parallel: Literal["cpu", "gpu"],
    num_cpu: int | None = None,
    gpus: list[int] = [],
) -> int | float:
    if parallel == "cpu":
        return test_cpu(x, num_parallel=num_cpu)
    elif parallel == "gpu":
        return test_gpu(x, gpus=gpus)

おわりに

まだまだユースケースはたくさんあります。
思い出したら続きを書きたいと思います。