pythonで@overloadを使って複数の引数の型に応じて返り値を変える

異なる型や引数を持つ関数を型ヒントを使って定義したい！

複数の異なる型や引数を持つ関数を型ヒント付きでしかもPythonで作ることは簡単ではありません。Union や TypeVar　の合併型や変数型では対応で対応できるものであればいいですが、そんなに世の中は甘くないようです。

では、そんな関数を定義するためにはどうしたらいいのか？

そこで、使われるのが、typing モジュールの @overload デコレータです！

実例

この例では、add 関数をオーバーロードしています。最初の2つのデコレータで異なる型の引数と戻り値を定義し、その後の通常の関数定義では、具体的な実装を提供しています。

この場合、add 関数は整数型の引数を受け取り整数を返すバージョンと、浮動小数点数型の引数を受け取り浮動小数点数を返すバージョンの2つが定義されています。適切な型の引数を渡すことで、適切なバージョンの関数が呼び出されます。

from typing import overload

@overload
def add(a: int, b: int) -> int:...

@overload
def add(a: float, b: float) -> float:...

def add(a, b):
    return a + b

以下は、この例を使用した実行例です。

print(add(1, 2))       # 整数の足し算: 3
print(add(1.5, 2.5))   # 浮動小数点数の足し算: 4.0

ただし、Pythonでは @overload デコレータによる静的な型チェックや実行時のオーバーロード解決は行われません。このデコレータは主に型チェッカーやIDEなどのツールが関数のシグネチャを認識しやすくするために使用されます。実行時のオーバーロード解決や型チェックは、通常の条件分岐や型の検査によって行われます。

@overload の注意点

@overload は、Pythonの実行時には無視されます。つまり、引数の型や数に応じた処理を実行するためには、実際にその関数を呼び出す必要があります。

また、@overload を使用する場合は、関数本体には処理を書いてはいけません。@overload は、あくまで関数のオーバーロードを定義するためのデコレーターであり、実際の処理はオーバーロードされた関数の本体に記述する必要があります。

より実践的なコード

typing モジュールの @overload デコレータを使用して、複数の引数の型に応じた挨拶文を生成する関数を定義しています。

from typing import overload

@overload
def generate_greeting(name: str) -> str:
    pass

@overload
def generate_greeting(name: str, age: int) -> str:
    pass

def generate_greeting(name, age=None):
    if age is None:
        return f"Hello, {name}!"
    else:
        return f"Hello, {name}! You are {age} years old."

# テストケース
print(generate_greeting("Alice"))   # "Hello, Alice!"
print(generate_greeting("Bob", 25)) # "Hello, Bob! You are 25 years old."

この例では、generate_greeting 関数をオーバーロードしています。最初のデコレータでは、name 引数のみを受け取り、文字列を返すバージョンを定義しています。2番目のデコレータでは、name 引数と age 引数を受け取り、文字列を返すバージョンを定義しています。

通常の関数定義の部分では、引数のデフォルト値を使用してオプションの age 引数を実現しています。引数の有無に応じて、適切な挨拶文を生成しています。

テストケースとして、generate_greeting 関数を呼び出して異なるパターンを試しています。適切なバージョンの関数が呼び出され、期待通りの挨拶文が生成されることを確認しています。

このように、@overload デコレータを使用することで、関数のオーバーロードを定義し、型ヒントを活用して異なる引数の組み合わせに対応する柔軟性を持った関数を作成できます。

@overloadとの出会い

そもそも、なんでこんな記事を書いたかというと、

機械学習系の論文や、パッケージのソースコードなんかを眺めているとよく、@overloadというデコレーターを最近になり見るようになりました。

なんのコードなんだかわからないけどあまり本筋に関係ないので、後回しにしていた理解のために今回はちゃんと記事にしてみました。

@overloadは、Pythonで関数の引数の型に応じて異なる処理を行いたい場合によく使われるようです。特に、引数の型や数が複雑な場合や、引数の組み合わせによって戻り値が変わるような関数での使用が一般的です。また、型ヒントを利用して、より明確かつ堅牢なコードを書くことができるようになるようです。