[Python] 関数, コルーチンどちらにも適用できるデコレータを実装する

概要

前々回と前回で普通の関数用とコルーチン用のretryデコレータの実装を紹介したので、今回はretryデコレータを普通の関数とコルーチンどちらに対しても適用できるようにします。

実装

asyncio.iscoroutinefunction()を使ってデコレートされる関数がコルーチンかどうかを判別して、普通の関数とコルーチン用のデコレータ実装を切り替えます。

_retry_sync(), _retry_async()の実装はほとんど以前の記事と同じです。

import asyncio
import functools
import time


def retry(tries: int = 1, retry_interval: float = 0,
          exceptions=(RuntimeError,)):
    assert tries > 0, "retries 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):
        if asyncio.iscoroutinefunction(func):
            return _retry_async(tries, retry_interval, exceptions)(func)
        return _retry_sync(tries, retry_interval, exceptions)(func)

    return decorator


def _retry_sync(tries: int = 1, retry_interval: float = 0,
                exceptions=(RuntimeError,)):
    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


def _retry_async(tries: int = 1, retry_interval: float = 0,
                 exceptions=(RuntimeError,)):
    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 %(name)s %(levelname)s %(message)s')


def f(threshold, logger):
    global count
    count += 1
    logger.info(f'{count}')

    if count < threshold:
        raise RuntimeError(count)


@retry(tries=3, retry_interval=1)
def f_sync(threshold):
    logger = logging.getLogger('f_sync')
    f(threshold, logger)


@retry(tries=3, retry_interval=1)
async def f_async(threshold):
    logger = logging.getLogger('f_async')
    f(threshold, logger)


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

        count = 0
        try:
            asyncio.run(f_async(value))
        except RuntimeError:
            logging.exception('f_async error')

動作確認

どちらでもリトライが動作していることが確認できました。

$ python3 main.py
2021-04-11 23:29:35,323 root INFO [threshold = 3]
2021-04-11 23:29:35,323 f_sync INFO 1
2021-04-11 23:29:36,327 f_sync INFO 2
2021-04-11 23:29:37,332 f_sync INFO 3
2021-04-11 23:29:37,332 f_async INFO 1
2021-04-11 23:29:38,338 f_async INFO 2
2021-04-11 23:29:39,341 f_async INFO 3
2021-04-11 23:29:39,342 root INFO [threshold = 4]
2021-04-11 23:29:39,342 f_sync INFO 1
2021-04-11 23:29:40,343 f_sync INFO 2
2021-04-11 23:29:41,346 f_sync INFO 3
2021-04-11 23:29:41,346 root ERROR f_sync error
Traceback (most recent call last):
  File "/Users/massakai/main.py", line 103, in <module>
    f_sync(value)
  File "/Users/massakai/main.py", line 29, in wrapper
    return func(*args, **kwargs)
  File "/Users/massakai/main.py", line 88, in f_sync
    f(threshold, logger)
  File "/Users/massakai/main.py", line 82, in f
    raise RuntimeError(count)
RuntimeError: 3
2021-04-11 23:29:41,347 f_async INFO 1
2021-04-11 23:29:42,352 f_async INFO 2
2021-04-11 23:29:43,354 f_async INFO 3
2021-04-11 23:29:43,355 root ERROR f_async error
Traceback (most recent call last):
  File "/Users/massakai/main.py", line 109, in <module>
    asyncio.run(f_async(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 52, in f
    return await func(*args, **kwargs)
  File "/Users/massakai/main.py", line 94, in f_async
    f(threshold, logger)
  File "/Users/massakai/main.py", line 82, in f
    raise RuntimeError(count)
RuntimeError: 3