日々精進

新しく学んだことを書き留めていきます

pythonでdecoratorを付けたメソッドの戻り値の型チェックができるようにする

以下のようにdecoratorを定義すると、@timerを付けたメソッドの戻り値の型がAnyになってしまい、不便。

def timer(fn) :
    from time import perf_counter

    def inner(*args, **kwargs):
        print(f'start {fn.__name__} {CommonUtil.now_datetime_for_log()}')
        CommonUtil.print_mem_usage()
        start_time = perf_counter()
        to_execute = fn(*args, **kwargs)
        end_time = perf_counter()
        execution_time = end_time - start_time
        hhmmss: str = time.strftime('%H:%M:%S', time.gmtime(execution_time))
        print(f'end {fn.__name__} {CommonUtil.now_datetime_for_log()}. elapsed time: {hhmmss}')
        CommonUtil.print_mem_usage()
        return to_execute

    return inner

def timer(fn) :

TCallable = TypeVar("TCallable", bound=Callable)

def timer(fn: TCallable) -> TCallable:

のように変えるとちゃんと戻り値の型をチェックしてくれ、コード補完もできる。

参考:

blog.whtsky.me