💪

Python で実行時型チェックをしよう:改良版

2023/12/13に公開

はじめに

よく調べたら前回作ったやつよりもシンプルかつ強力なのができたので共有します。

Python のバージョンは 3.10 以上にしてください。

参考文献

Install

$ pip install nobus
from nobus.typecheck import typecheck

コピペ用

ロクにメンテされていない個人開発のライブラリに抵抗がある人は以下のソースコードをコピペしてください。50行もなくとても素直で読めば分かると思いますので解説もしません。

デコレータとしての威力は functools.wraps によるもので、型アノテーションの扱いは inspect.signature によるものです。

# MIT License
# Copyright (c) 2023 Yoshinobu Ogura

from inspect import signature, isfunction, ismethod
from functools import wraps

def typecheck(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        def _typecheck(arg, annot, v):
            if isfunction(annot) or ismethod(annot):
                # 型アノテーションが関数やメソッドで与えられているとき
                if not annot(arg):
                    raise TypeError(f"{v} does not satisfy condition {annot.__name__}.")
            elif not isinstance(arg, annot):
                # それ以外は isinstance で処理
                raise TypeError(f"{v} must be instance of {annot} but actual type {type(arg)}.")
        
        # 型シグネチャを取得
        sgnt = signature(f)
        # 与えられた引数をパラメータにバインド
        bind = sgnt.bind(*args, **kwargs)
        # デフォルト引数をバインド
        bind.apply_defaults()

        for p, arg in bind.arguments.items():
            # パラメータに対応する型アノテーションを取得
            param = sgnt.parameters[p]
            annot = param.annotation
            
            if annot is param.empty:
                # 型アノテーションがなければ無視
                continue

            _typecheck(arg, annot, v=p)

        # 関数を実行
        ret = f(*args, **kwargs)

        # 戻り値の型アノテーションを取得
        annot = sgnt.return_annotation
        if annot is not sgnt.empty:
            # 戻り値に型アノテーションが存在するとき
            _typecheck(ret, annot, v="return value")
        
        return ret

    return wrapper

Usage

typecheck はデコレータで、与えられた関数の型アノテーションを読んで isinstance による実行時型チェックを行います。型アノテーションに関数またはメソッドを与えた場合は、その関数にアノテートされている引数を与え、結果が False となるとき型エラーになります。

アノテートしないとき

何も起こりません。

@typecheck
def f(x, y):
    return x + y

f(1, 2)
#>>> 3

関数の引数に型アノテーションをつけるとき

isinstance を実行して False になるときエラーになります。

関数の定義
@typecheck
def f(x, y: int):
    return x + y

# isinstance(2, int) == True なので OK
f(1, 2)
#>>> 3

# isinstance('b', int) == False なのでエラー
f(1, 'b')
#>>> エラー
表示されるエラー
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[14], line 1
----> 1 f(1, 'b')
...

TypeError: y must be instance of <class 'int'> but actual type <class 'str'>.

関数の戻り値に型アノテーションをつけるとき

同様に、isinstance を実行して False になるときエラーになります。

@typecheck
def f(x, y) -> int:
    return x + y

# isinstance(x + y, int) == True なので OK
f(1, 2)
#>>> 3

# isinstance(x + y, int) == False なのでエラー
f('a', 'b')
#>>> エラー
表示されるエラー
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[17], line 2
      1 # isinstance(x + y, int) == False なのでエラー
----> 2 f('a', 'b')
...

TypeError: return value must be instance of <class 'int'> but actual type <class 'str'>.

関数による条件付与を行う特殊なアノテーション

Python の型アノテーションは意外とヤンチャできるので、以下のように引数に条件を付与することができます。条件付与を行う関数の型は Callable[[Any], bool] です。条件付与関数の実行結果が False のとき TypeError になります。

def odd(x):
    return x % 2 == 1

def even(x):
    return x % 2 == 0

@typecheck
def f(x: odd, y: odd) -> even:
    return x + y

# odd(x) == True, odd(y) == True,
# even(x + y) == True なので OK
f(1, 3)
#>>> 4

# odd(x) == False なのでエラー
f(2, 3)
#>>> エラー
表示されるエラー
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[18], line 13
     12 # odd(x) == False なのでエラー
---> 13 f(2, 3)
...

TypeError: x does not satisfy condition odd.

クラスのメソッドでアノテーションを与えることも可能です。classmethodstaticmethod にも対応しており、かなり強力です。

class A:
    def odd(self, x):
        return x % 2 == 1

    @classmethod
    def even(cls, x):
        return x % 2 == 0

    @staticmethod
    def fizz(x):
        return x % 3 == 0

a = A()

@typecheck
def f(x: a.odd, y: A.even) -> A.fizz:
    return x + y

# 以下を試してみてください
f(1, 2)
#>>> 3

f(2, 2)
#>>> x がエラー

f(1, 3)
#>>> y がエラー

f(3, 4)
#>>> 戻り値がエラー

メソッドにアノテーションするとき

メソッドのときも使い方は同様です。とても強力です。

class A:
    @typecheck
    def f(self, x: int, y: int):
        return x + y

    @classmethod
    @typecheck
    def g(cls, x: int, y: int):
        return x + y

    @staticmethod
    @typecheck
    def h(x: int, y: int):
        return x + y

# 与える引数を変えて試してみてください
a = A()
a.f(1, 2)
#>>> 3

A.g(1, 2)
#>>> 3

A.h(1, 2)
#>>> 3

複雑な関数定義にアノテーションするとき

positional only, keyword only, デフォルト引数など、諸々の関数定義についてもすべて対応しています。遠慮なく酷使してください。可変長は試してないのでどうなるか知りません。

@typecheck
def f(x: int, y: int, /, z: int, *, a: int, b: int=2):
    return a * x + b * y + z

f(1, 2, z=3, a=4)

Discussion