株式会社コスモソニックツーワンの情報発信ブログです。

Pythonで関数の引数型チェックをデコレータで実装する

あけましておめでとうございます。サカイリです。
新年あけて初めての投稿はPythonにて自作関数の型チェックをデコレータとアノテーションを利用して実装してみました。
(この記事は私がQiitaで投稿した同名記事『Pythonで関数の引数型チェックをデコレータで実装する』と同じ内容となっています。)



概要

自作関数に引数の型チェックを実装する際にデコレータとアノテーションを利用して実装する。
これがあればデコレートするだけで型チェックができるね!

偉大なる先人

Python関数の引数と返り値の型をチェックするデコレータ
↑やりたかったことはほぼこれです。
python3.6系でエラーになっていたので一部修正とアノテーション周りの整備をしてみました。
また、返り値のチェックは不要だったので削除しています。

デコレータの実装

今回の型チェックの実装は以下とします。

  • 関数の引数について型チェックを行う
  • アノテーションで型が指定されている引数をチェック対象とする
  • アノテーションがコメントだったり、ない場合はスルーする
import inspect
import functools

def args_type_check(func):
    @functools.wraps(func)
    def args_type_check_wrapper(*args, **kwargs):
        """
        引数についてアノテーションで指定した型と一致しているかのチェックを行うデコレータ。
        """
        sig = inspect.signature(func)
        for arg_key, arg_val in sig.bind(*args, **kwargs).arguments.items():
            annot = sig.parameters[arg_key].annotation
            request_type = annot if type(annot) is type else inspect._empty
            if request_type is not inspect._empty and type(arg_val) is not request_type:
                error_msg = '引数"{}"の型が対応していません。(対応している型:{}、指定された型:{})'
                raise TypeError(error_msg.format(arg_key, request_type, type(arg_val)))
        return func(*args, **kwargs)
    return args_type_check_wrapper

細かくみてみましょう。

@functools.wraps(func)

これがないと関数名やDocstringが自作関数でなくデコレータのものが表示されます。

request_type = annot if type(annot) is type else inspect._empty

アノテーションは型以外にも文字列の場合があるので、型の時だけ取り出すようにします。

自作関数へのデコレータ適用

実際に関数にデコレータを適用するときはこのようになります。

@args_type_check
def test_func(arg1:str, arg2:str='arg2', arg3:'this is arg3.'='arg3', arg4='arg4'):
    print(arg1, arg2, arg3, arg4)

型チェックが行われるのはアノテーションにて型を指定した引数のみになるため、
`arg1`と`arg2`については型チェックが行われ`arg3`と`arg4`はスルーされます。

終わりに

リストや辞書型の場合の中身についてもチェックが必要な場合も実施できるようにする予定です。また作業時間ができたらこちらに反映します。