Pythonで async def / def 両対応のデコレータを書く

タイトルの通り。

単純に以下の様なdecoratorを書くと、async defをラップできない。

from functools import wraps
from datetime import datetime

def timetrack(func):
    @wraps(func)
    def inner(self, *args, **kwargs):
        start = datetime.now()
        return_value = func(self, *args, **kwargs)
        end = datetime.now()
        self.logger.info('%s takes %s sec', func.__name__, (end - start).total_seconds())
        return return_value
    return inner

かといって次のように書くと、今度は普通のdefが勝手にcoroutine functionになってしまって困る。

from functools import wraps
from datetime import datetime

def timetrack(func):
    @wraps(func)
    async def inner(self, *args, **kwargs):
        start = datetime.now()
        return_value = await func(self, *args, **kwargs)
        end = datetime.now()
        self.logger.info('%s takes %s sec', func.__name__, (end - start).total_seconds())
        return return_value
    return inner

なので、asyncio.iscoroutinefunctionを使って以下のように分岐するといける。

from functools import wraps
from datetime import datetime

def timetrack(func):
    if asyncio.iscoroutinefunction(func):
        @wraps(func)
        async def async_inner(self, *args, **kwargs):
            start = datetime.now()
            return_value = await func(self, *args, **kwargs)
            end = datetime.now()
            self.logger.info('%s takes %s sec', func.__name__, (end - start).total_seconds())
            return return_value
        return async_inner
    else:
        @wraps(func)
        def inner(self, *args, **kwargs):
            start = datetime.now()
            return_value = func(self, *args, **kwargs)
            end = datetime.now()
            self.logger.info('%s takes %s sec', func.__name__, (end - start).total_seconds())
            return return_value
        return inner

一般化したユーティリティみたいのを書けないこともない気がするけど、今回そこまでのモチベーションはなし。