feat: timing decorator supports sync and async functions

Author: NightMarcher

Pipfile

@@ -14,6 +14,7 @@ cryptography = "*"
 faker = "*"
 fastapi = "*"
 pytest = "*"
+pytest-asyncio = "*"
 tortoise-orm = "*"
 uvicorn = "*"
 

fastapi_esql/utils/decorator.py

@@ -1,3 +1,4 @@
+import asyncio
 from functools import wraps
 from logging import getLogger
 from time import perf_counter
@@ -6,11 +7,24 @@
 
 
 def timing(func):
+
     @wraps(func)
-    async def wrapper(*args, **kwargs):
+    async def async_wrapper(*args, **kwargs):
         st = perf_counter()
         result = await func(*args, **kwargs)
         et = perf_counter()
         logger.info(f"{func.__module__}/{func.__name__} Cost {1000 * (et - st):.2f} ms")
         return result
-    return wrapper
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        st = perf_counter()
+        result = func(*args, **kwargs)
+        et = perf_counter()
+        logger.info(f"{func.__module__}/{func.__name__} Cost {1000 * (et - st):.2f} ms")
+        return result
+
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    else:
+        return wrapper

requirements-dev.txt

@@ -5,5 +5,6 @@ cryptography
 faker
 fastapi
 pytest
+pytest-asyncio
 tortoise-orm
 uvicorn

tests/test_decorator.py

@@ -0,0 +1,26 @@
+import pytest
+
+from fastapi_esql import timing
+
+
+@timing
+async def afunc():
+    """async def afunc()"""
+    return True
+
+
+@timing
+def func():
+    """def func()"""
+    return False
+
+
+@pytest.mark.asyncio
+async def test_async_timing():
+    print(afunc, afunc.__doc__, afunc.__name__)
+    assert await afunc() == True
+
+
+def test_sync_timing():
+    print(func, func.__doc__, func.__name__)
+    assert func() == False