[Python] typing.overload で引数に応じて返却値の型を変える

typing.overload を使うと、引数に応じて、返却値の型を定義できる。

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

例えば、キャッシュデータがあり、そこからキーを指定してデータを取り出すような関数。

from enum import Enum
from models.user import User

cached = {}

class CacheKey(Enum):
    # 現在のユーザー数
    USER_COUNT = 'USER_COUNT'
    # 直近登録されたユーザー
    LATEST_USER = 'LATEST_USER'

def get_by_key(key: CacheKey) -> int | User | None:
    return getattr(cached, key.name)

typing.overload を使うとこうなる。

from enum import Enum
from typing import overload, Literal
from models.user import User

cached = {}

class CacheKey(Enum):
    # 現在のユーザー数
    USER_COUNT = 'USER_COUNT'
    # 直近登録されたユーザー
    LATEST_USER = 'LATEST_USER'


@overload
def get_by_key(key: Literal[CacheKey.USER_COUNT]) -> int | None: ...

@overload
def get_by_key(key: Literal[CacheKey.LATEST_USER]) -> User | None: ...

@overload
def get_by_key(key: CacheKey) -> any: ...

def get_by_key(key: CacheKey):
    return getattr(cached, key.name)

これで、キーに応じた型がつく。

count = get_by_key(CacheKey.USER_COUNT) # int | None
user = get_by_key(CacheKey.LATEST_USER) # User | None