🎀

PythonのDecoratorを利用した引数・戻り値の型チェック実装

2023/12/05に公開

概要

自作関数を利用する際に、引数と戻り値の型チェックを行うことで、コードの安全性を高められます。
Pythonのデコレータとアノテーションを利用することで、簡単に複数関数へ型チェックを適用することができます。

デコレータとは

デコレータは関数やクラスをラップして、別の関数を返す関数のことです。ここでは関数デコレータについて話を進めていきます。
既存の関数をラップすることにより、引数に指定した関数の前後で処理を追加することができます
例えば「引数を処理してから関数を呼び出す」や「返り値に処理を追加して変更する」などがあげられます。

デコレータの使い方

Pythonではラップする関数の宣言時に@decorator関数名を追加することで簡単に適用することができます。@classmethod@staticmethodが一般的によく見るデコレータです。
以下の例では関数「say_hello_world」にデコレータ関数「example_decorator」を追加することで前後処理を追加してラップしています。

# デコレータ関数
def example_decorator(func):
    def example_wrapper():
        print('前処理の追加')
        func()
        print('後処理の追加')
    return example_wrapper

# デコレートされた関数
@example_decorator
def say_hello_world():
    print('Hello world!')

ここで便宜上デコレータ関数と読んでいますが、「example_decorator」は関数を返す関数です。そのため、以下のような呼び出しでも同様の結果が得られます。

wraped_func = example_decorator(say_hello_world)

アノテーションとは

アノテーションは関数やクラスに対して付与できるの型の注釈のことです。
引数や戻り値、クラスの属性などについてそのデータ型を明示することができますが、あくまで注釈なので型チェックが行われるわけではありません
そのため、アノテーションで指定した型以外の型を渡してもエラーなどにはなりません。

引数・戻り値の型チェックをデコレータで実装

それではアノテーションを利用した引数・戻り値の型チェックデコレータ関数を実装します。
今回の要件は以下の通りです。

  • 引数・戻り値の型チェックを行う
    • アノテーションで型が指定されているものをチェック対象とする
    • アノテーションがコメントだったり、存在しない場合はスルーする
    • 指定された型ではない場合はTypeError例外を発生させる
def type_check(func):
    @functools.wraps(func)
    def type_check_wrapper(*args, **kwargs):
        '''引数・戻り値の型チェックを行うデコレータ。
        '''
        func_signature = inspect.signature(func)

        # 引数の型チェック
        args_dict = func_signature.bind(*args, **kwargs)
        for arg_key, arg_value in args_dict.arguments.items():
                # 引数のアノテーションを取得
                arg_annotation = func_signature.parameters[arg_key].annotation
                specified_arg_type = arg_annotation if type(arg_annotation) is type else inspect._empty
                # 引数とアノテーションの型を比較
                if specified_arg_type is not inspect._empty and type(arg_value) is not specified_arg_type:
                    error_msg = '引数"{}"の型が対応していません。(アノテーション:{}、引数の型:{})'
                    raise TypeError(error_msg.format(arg_key, specified_arg_type, type(arg_value)))

        # 関数の実行
        results = func(*args, **kwargs)

        # 戻り値の型チェック
        return_annotation = func_signature.return_annotation
        specified_return_type = return_annotation if type(return_annotation) is type else inspect._empty
        # 戻り値とアノテーションの型を比較
        if specified_return_type is not inspect._empty and type(results) is not specified_return_type:
            error_msg = '戻り値の型が対応していません。(アノテーション:{}、戻り値の型:{})'
            raise TypeError(error_msg.format(specified_return_type, type(results)))

        return results
    return type_check_wrapper

いくつか補足します。

@functools.wraps(func)

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

specified_arg_type = arg_annotation if type(arg_annotation) is type else inspect._empty
specified_return_type = return_annotation if type(return_annotation) is type else inspect._empty

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

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

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

@type_check
def test_func(
    arg1:str,
    arg2:str='arg2',
    arg3:'this is arg3'='arg3',
    arg4='arg4'
) -> str:
    return str(arg1) + str(arg2) + str(arg3) + str(arg4)

print(test_func('t1','t2', 't3', 't4'))

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

課題

typing._GenericAliasの考慮と複数の戻り値を設定した場合の処理を追加する。
(ただ、複数の戻り値を返す関数は、挙動が変わった時に予期せぬエラーの原因になりやすいので、あんまり使わない方が良いと私は思う。もしも必要な場合は独自クラス作るか、オブジェクトで返してあげるのが良さそうです。)

参考

Python関数の引数と返り値の型をチェックするデコレータ
↑偉大なる先人がいらっしゃいました。
python3.12系でエラーになっていたので修正とアノテーション周りの処理を整備をしています。

Aidemy Tech Blog

Discussion