[Python] コルーチンに適用するデコレータを実装する (retryデコレータ)

概要

retryデコレータの実装を通してコルーチンに適用するデコレータの実装方法を紹介します。

背景

ロギングやメトリクスの計測などの実装をロジックとは分離したいので、デコレータとして実装します。

Python3ではasync defでコルーチンを定義することができますが、コルーチンに適用可能なデコレータは通常の関数に対するデコレータとは実装方法が若干異なります。

実装

デコレートするコルーチンはコルーチンを返すようにするところがポイントです。

import asyncio
import functools


def retry(tries: int = 1, retry_interval: float = 0,
          exceptions=(RuntimeError,)):
    assert tries > 0, "retries must be greater than or equal to 0."
    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):
            async def f():
                num = tries
                while True:
                    try:
                        return await func(*args, **kwargs)
                    except Exception as e:
                        num -= 1
                        if (not any(isinstance(e, exception)
                                    for exception in exceptions)
                                or num <= 0):
                            raise
                    await asyncio.sleep(retry_interval)

            return f()

        return wrapper

    return decorator


# ここから下は確認用のコード
import logging

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)-15s %(levelname)s %(message)s')


@retry(tries=3, retry_interval=1)
async def f(threshold):
    global count
    count += 1
    logging.info(f'{count}')

    if count < threshold:
        raise RuntimeError(count)


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

動作確認

大体1秒間隔でリトライしていることがわかります。

[massakai@MacBook-Pro-15-inch-2017]% python3 main.py
2021-04-09 09:50:09,467 INFO [threshold = 3]
2021-04-09 09:50:09,467 INFO 1
2021-04-09 09:50:10,472 INFO 2
2021-04-09 09:50:11,476 INFO 3
2021-04-09 09:50:11,478 INFO [threshold = 4]
2021-04-09 09:50:11,478 INFO 1
2021-04-09 09:50:12,483 INFO 2
2021-04-09 09:50:13,486 INFO 3
Traceback (most recent call last):
  File "/Users/massakai/main.py", line 58, in 
    asyncio.run(f(value))
  File "/usr/local/Cellar/python@3.9/3.9.1_8/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/usr/local/Cellar/python@3.9/3.9.1_8/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/base_events.py", line 642, in run_until_complete
    return future.result()
  File "/Users/massakai/main.py", line 19, in f
    return await func(*args, **kwargs)
  File "/Users/massakai/main.py", line 50, in f
    raise RuntimeError(count)
RuntimeError: 3

蛇足

元々はTornadoの非同期ハンドラのレイテンシを計測するデコレータを実装する予定だったのですが、前に記事にした通りon_finish()でシンプルに実装できてしまいました。そのため、この記事ではリトライを題材にしています。