[Python] パラメータ付きのデコレータを実装する (retryデコレータ)

概要

パラメータ付きのデコレータの実装例として、retryデコレータを実装しました。

実装

パラメータ付きのデコレータは、パラメータなしのデコレータ(関数1つを引数にとる関数)を返す関数として定義します。

import functools
import time


def retry(tries: int = 1, retry_interval: float = 0,
          exceptions=(RuntimeError,)):
    assert tries > 0, "tries must be positive integer."
    assert retry_interval >= 0, (
        "retry_interval must be greater than or equal to 0.")
    assert exceptions, "exceptions must not be empty."

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            num = tries
            while True:
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    num -= 1
                    if (not any(isinstance(e, exception)
                                for exception in exceptions)
                            or num <= 0):
                        raise
                time.sleep(retry_interval)

        return wrapper

    return decorator

動作確認

適当に書いた以下のコードで動作を確認します。

@retry(tries=3, retry_interval=0.1)
def f(threshold):
    global count
    count += 1
    print(f'{count}', flush=True)

    if count < threshold:
        raise RuntimeError(count)


if __name__ == '__main__':
    # 3回試行して成功するケースと3回試行して失敗するケース
    for value in (3, 4):
        count = 0
        print(f"[threshold = {value}]")
        f(value)

実行すると2回までリトライしていることがわかります。

$ python3 main.py
[threshold = 3]
1
2
3
[threshold = 4]
1
2
3
Traceback (most recent call last):
  File "/Users/massakai/main.py", line 47, in 
    f(value)
  File "/Users/massakai/main.py", line 18, in wrapper
    return func(*args, **kwargs)
  File "/Users/massakai/main.py", line 39, in f
    raise RuntimeError(count)
RuntimeError: 3