この記事は Acompany5周年アドベントカレンダー 8日目 の記事です．
はじめに
動的型付け言語であるPythonは関数のオーバーロードを特殊な処理で実現しています．過去の記事では，標準で提供されているオーバーロードを実現する関数であるsingledispatchの仕組みを説明し，型の代替となる値を自分で定義することによってsingledispatchの拡張を行いました．
ここで，Pythonには型アノテーションと呼ばれる型付けが可能であることに着目します．過去では型情報が取得できないという仮定の元で様々な拡張をしていましたが，この型アノテーション情報を取得して活用できればより良いオーバーロードが実現できそうです．実現します．
singledispatchの仕組み
まずはPython標準のsingledispatchがどのように動いているかについて簡単に復習します．
singledispatchは主にregistry，register()，wrapper()という3つの要素から構成されており，それぞれ次の役割を持ちます．
registry- オーバーロードする関数を格納する辞書型変数
{type: func}という形式で保存する
register()- オーバーロードする関数を
registryに保存する関数
- オーバーロードする関数を
wrapper()- 関数呼び出し時に呼ばれて引数に応じて呼ぶ関数を切り替える関数
まず，オーバーロードしたい関数にsingledispatchをデコレートしてその関数の代わりにwrapper()が呼ばれるようにします．次に，オーバーロードで呼びたい関数にregister()をデコレートして関数をregistryに登録します．これでオーバーロードが実現できます．
以下に最小限の要素で構成した自作dispatchによるオーバーロードと簡単なイメージ図を示します．
# 最小限の機能だけ携えた自作のdispatch def mydispatch(func): registry = {} def wrapper(*args, **kw): return registry[args[0].__class__](*args, **kw) def register(cls, func=None): if func is None: return lambda f: register(cls, f) registry.setdefault(cls, func) return func wrapper.register = register return wrapper @mydispatch def f(*args, **kw): raise NotImplementedError("") @f.register(int) def f_int(n: int) -> None: print("int") @f.register(str) def f_str(s: str) -> None: print("str") if __name__ == "__main__": f(5) f("str")
int str
型アノテーション情報の取得
さて，オーバーロードを実現するsingledispatchの最低限の構造が分かったので早速型アノテーションを用いた拡張を考えていきます．型アノテーションはmypyなどの型チェッカに用いられるもので，実行時には無視されます．そのため実行時には型情報は取得できないように見えますが，実は取得する手段が存在します．ここでは関数に対する型アノテーションの取得方法を紹介します．
__annotations__
型アノテーション取得方法の1つがmagic methodの1つである__annotations__です．__annotations__は引数の型，返り値の型を{変数名: 型}という形式の辞書型で保存しています．ただし，返り値の型の名前はreturnとなります．以下に__annotations__の使用例を示します．
class C: def f(self, a: int, b: str) -> float: return 0.0 def f(a: int, b: str) -> float: return 0.0 if __name__ == "__main__": print(f.__annotations__) print(C().f.__annotations__)
{'a': <class 'int'>, 'b': <class 'str'>, 'return': <class 'float'>}
{'a': <class 'int'>, 'b': <class 'str'>, 'return': <class 'float'>}
get_type_hints()
get_type_hints()も型アノテーション取得方法の1つで，内部で__annotations__を呼んでいるため基本的には同じ挙動になります．以下にget_type_hints()の使用例を示します．
from typing import get_type_hints # 関数定義省略 if __name__ == "__main__": print(get_type_hints(f)) print(get_type_hints(C().f))
{'a': <class 'int'>, 'b': <class 'str'>, 'return': <class 'float'>}
{'a': <class 'int'>, 'b': <class 'str'>, 'return': <class 'float'>}このように__annotations__と同等の結果が得られました．
get_type_hints()の__annotations__より便利な要素として型の遅延評価が可能な点があります．例えば，あるクラス内でそのクラスを型としてアノテーションする場合，アノテーションの時点ではクラスが定義されていないため文字列としてアノテーションすることになります．これが，__annotations__では文字列としてしか解釈できないのですが，get_type_hints()を用いると実行時に(デフォルトだと)スコープ内から適切な型を探索して置き換えます．以下に使用例を示します．
from typing import get_type_hints class C: def f(self, a: int, b: "str") -> "C": C() if __name__ == "__main__": print(C().f.__annotations__) print(get_type_hints(C().f))
{'a': <class 'int'>, 'b': 'str', 'return': 'C'}
{'a': <class 'int'>, 'b': <class 'str'>, 'return': <class '__main__.C'>}__annotations__では型が文字列で表現されていますが，get_type_hints()は型として解釈されていることが分かります．基本的にはget_type_hints()の方が便利なため，以下ではこれを使用していきます．
型アノテーション情報を用いたオーバーロード
さて，必要な情報は揃ったので実際に型アノテーション情報を用いたオーバーロード機構を構築します．
1つの引数についてオーバーロード
先ほど紹介した最小限のmydispatchをベースとして，まずはregisterする情報を型アノテーション情報に変更します．
def mydispatch(func): # 中略 def register(func): _, cls = next(iter(get_type_hints(func).items())) registry[cls] = func return func # 後略
これだけで変更は終了です．簡単ですね．実際の使用例を以下に示します．
@mydispatch def f(*args, **kw): print("NotImplementedError") @f.register def f_int(n: int) -> None: print(f"int: {n}") @f.register def f_str(s: str): print(f"str: {s}") if __name__ == "__main__": f(10) f("Hello")
int: 10 str: Hello
明示的に型を指定することなく，型アノテーション情報だけでオーバーロードできています．なお，Python標準のsingledispatchでもここまでは実現されています．
複数の引数についてオーバーロード
上で示したコードは1つ目の引数のみに基づいて判定がなされています．実際の使用時は複数の引数でオーバーロードしたいと思うので拡張します．get_type_hintsで全ての引数の型は取得できるので，これをそのままtupleにしてkeyにします．ただし，返り値の型は含めないように注意します．
def mydispatch(func): # 中略 def register(func): key = [ cls for argname, cls in get_type_hints(func).items() if argname != "return" ] registry[tuple(key)] = func return func def wrapper(*args, **kw): key = tuple([x.__class__ for x in args]) return registry[key](*args, **kw) # 後略
使用例を以下に示します．
@mydispatch def f(*args, **kw): print("NotImplementedError") @f.register def f_int(n: int, a: int) -> None: print(f"int: {n}, a: {a}") @f.register def f_str(n: int, s: str): print(f"int: {n}, str: {s}") if __name__ == "__main__": f(10, 10) f(10, "Hello")
int: 10, a: 10 int: 10, str: Hello
2つ目の引数についても正しくオーバーロードできました．
Generics型についてオーバーロード
PythonはList[int]やTuple[str]のようなGeneric型も扱うことができます．現在型の判定で用いている__class__ではGeneric型を判定できないので，この部分を自作してGeneric型も扱えるように拡張します．
なお，都合上Python3.8以前を仮定します．Python3.9以降では非推奨な機能が使用されていますが，使う場合は適宜置き換えてください．
from typing import List, get_type_hints def __get_type(val): t = type(val) if t == list: return typing.List[__get_type(val[0])] return t def mydispatch(func): # 中略 def wrapper(*args, **kw): key = tuple([__get_type(x) for x in args]) return registry[key](*args, **kw) # 後略
使用例を以下に示します．
@mydispatch def f(*args, **kw): print("NotImplementedError") @f.register def f_int_int(a: int, b: int) -> None: print(f"int int: {a}, {b}") @f.register def f_list_int(a: List[int], b: int) -> None: print(f"list int: {a} {b}") @f.register def f_int_listlist(a: int, b: List[List[int]]) -> None: print(f"int list[list]: {a} {b}") if __name__ == "__main__": f(1, 1) f([1], 1) f(1, [[1]])
int int: 1, 1 list int: [1] 1 int list[list]: 1 [[1]]
Generics型についても正しくオーバーロードできました．今回は簡素にするためListのみ追加していますが，よしなに追加すればどの型でも可能です．
余談ですが，あらゆる型についてオーバーロードを実現するには，あらゆる値からtypingの型へと完全に変換することが必要です．これは実質型チェックをするのと同様です．つまり，完璧なオーバーロードをするためには型チェッカを自前実装する必要があります．えぇ...
メソッドもオーバーロードする
クラス内に定義するメソッドについても同様にしてオーバーロードできそうなものですが，selfの存在によって正常に動作しません．メソッドの引数の型をget_type_hintsで取得すると第一引数のselfが除去された辞書が生成されますが，関数呼び出し時にはkeyの中にselfが入ってしまいオーバーロードに失敗します．
そこで，メソッドの時だけ第一引数を無視する分岐を加えることにします．この実現のためにまずは普通の関数とメソッドとを区別する方法を考えます．結論この方法は分かりませんでした．そのため，以下に示す方法で強引に区別します．
def __is_method_probably(func): return "." in func.__qualname__ class C: def f(self, n: int): ... def f(n: int): ... if __name__ == "__main__": print(__is_method_probably(f)) print(__is_method_probably(C().f))
False True
普通の関数とメソッドの区別ができました．正しく区別できる保証はないのでprobablyとしています．より良い方法があればご教授ください．
さて，ではメソッドと関数とで処理を分岐させます．
def __is_method_probably(func): return "." in func.__qualname__ def mydispatch(func): is_method = __is_method_probably(func) # 中略 def wrapper(*args, **kw): key = tuple([__get_type(x) for x in (args[1:] if is_method else args)]) return registry[key](*args, **kw) # 後略
使用例を以下に示します．
class C: @mydispatch def f(*args, **kw): print("NotImplementedError") @f.register def f_int(self, n: int): print("int") @f.register def f_list(self, n: List[int]): print("list") @mydispatch def f(*args, **kw): print("NotImplementedError") @f.register def f_int(n: int): print("int") @f.register def f_list(n: List[int]): print("list") if __name__ == "__main__": c = C() c.f(1) c.f([1]) f(1) f([1])
int list int list
関数もメソッドも正しくオーバーロードできました．
完成系
import typing from functools import update_wrapper from typing import List, get_type_hints def __get_type(val): t = type(val) if t == list: return typing.List[__get_type(val[0])] return t def __is_method_probably(func): return "." in func.__qualname__ def mydispatch(func): registry = {} is_method = __is_method_probably(func) def register(func): key = [ cls for argname, cls in get_type_hints(func).items() if argname != "return" ] registry[tuple(key)] = func return func def wrapper(*args, **kw): key = tuple([__get_type(x) for x in (args[1:] if is_method else args)]) return registry[key](*args, **kw) wrapper.register = register update_wrapper(wrapper, func) return wrapper
