何かを書き留める何か

数学や読んだ本について書く何かです。最近は社会人として生き残りの術を学ぶ日々です。

Pythonで函数の型チェック

先日、 第6回Python文法詳解を詳解する会に参加した。 python-in-depth.connpass.com

デコレータと聞いたので、予習としてデコレータによる型チェック機能をPythonで実装した。 引数とアノテーションに関する情報はinspectモジュールで取得、逐一isinstance函数で引数をチェックし、最後に元の返り値をisinstance函数でチェック、という流れ。 デフォルト値とは異なるキーワード専用引数はinspectで拾えない(と思う)ので素直にwrapper函数の定義の際に**kwargsで拾うことに。

実は Mark Summerfield『Python 3 プログラミング徹底入門』という本に実装が書いてあり、結局それを眺めることになったが、Summerfield氏の実装よりもitertoolsモジュールを使うことでちょっとシンプル(かつ多分効率的?)になったと思いたい。

アノテーションfloat型の引数にint型が渡された場合、適当にキャストして欲しい場合が多いと思われるが、この実装だと融通が利かずTypeErrorを返してしまう。

import functools
import inspect
import itertools


def typecheck(func):
    """
    型チェック用デコレータ

    関数のアノテーションに基づいて引数及び返り値の型をチェックする。
    型がアノテーションと一致しない場合はTypeErrorを投げる
    """

    arg_spec = inspect.getfullargspec(func)

    @functools.wraps(func)
    def wrapper(*args, **kwags):
        """
        ラッパ関数

        前半部分で引数の型チェック、後半で返り値の型チェックを行う
        """
        for name, arg in itertools.chain(zip(arg_spec.args, args), kwags.items()):
            if not isinstance(arg, arg_spec.annotations[name]):
                raise TypeError("Arg {0} is expected {1}, but got {2}".format(
                    name, arg_spec.annotations[name], type(arg)))
        result = func(*args, **kwags)
        if not isinstance(result, arg_spec.annotations['return']):
            raise TypeError("Return Value is expected {0}, but got {1}".format(
                arg_spec.annotations['return'], type(result)))
        return result
    return wrapper

if __name__ == '__main__':
    @typecheck
    def intsum(a: int, b: int, multiple: int=1, *, divide: int=1) -> int:
        """
        int型の足し算を行う

        オプションによって掛け算や割り算も行う
        """
        return ((a + b) * multiple) // divide

    print(intsum(1, 2))
    print(intsum(3, 2, multiple=3))
    print(intsum(5, 6, divide=2))
    try:
        print(intsum(3, 2, '1'))
    except TypeError as e:
        print(e)