Add support for decorating async functions (#33)

* Add support for decorating async functions

* Update types for async

* Add a test of a basic async function

* Fix typo in fastapi-example.py

* Format code

* Update fastapi-example.py with async code as well

* Implement PR feedback (fix docstrings, rename symbol)

* Add test for async exceptions

* Update django dep in example to resolve dependabot alert

---------

Co-authored-by: Brett Beutell <brett@fiberplane.com>

Author: brettimus

examples/django_example/django_example/asgi.py

@@ -11,6 +11,6 @@
 
 from django.core.asgi import get_asgi_application
 
-os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'django_example.settings')
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "django_example.settings")
 
 application = get_asgi_application()

examples/django_example/django_example/wsgi.py

@@ -11,6 +11,6 @@
 
 from django.core.wsgi import get_wsgi_application
 
-os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'django_example.settings')
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "django_example.settings")
 
 application = get_wsgi_application()

examples/django_example/manage.py

@@ -6,7 +6,7 @@
 
 def main():
     """Run administrative tasks."""
-    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'django_example.settings')
+    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "django_example.settings")
     try:
         from django.core.management import execute_from_command_line
     except ImportError as exc:
@@ -18,5 +18,5 @@ def main():
     execute_from_command_line(sys.argv)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

examples/example.py

@@ -94,6 +94,6 @@ def random_error():
 
     try:
         # Call random_error. It will randomly raise an error or return "ok"
-        random_error() 
+        random_error()
     except Exception:
         pass

examples/fastapi-example.py

@@ -1,3 +1,4 @@
+import asyncio
 from fastapi import FastAPI, Response
 import uvicorn
 from autometrics import autometrics
@@ -7,24 +8,40 @@
 
 
 # Set up a metrics endpoint for Prometheus to scrape
-# `generate_lates` returns the latest metrics data in the Prometheus text format
+# `generate_latest` returns the latest metrics data in the Prometheus text format
 @app.get("/metrics")
 def metrics():
     return Response(generate_latest())
 
 
 # Set up the root endpoint of the API
-@autometrics
 @app.get("/")
+@autometrics
 def read_root():
     do_something()
     return {"Hello": "World"}
 
 
+# Set up an async handler
+@app.get("/async")
+@autometrics
+async def async_route():
+    message = await do_something_async()
+    return {"Hello": message}
+
+
 @autometrics
 def do_something():
     print("done")
 
 
+@autometrics
+async def do_something_async():
+    print("async start")
+    await asyncio.sleep(2.0)
+    print("async done")
+    return "async world"
+
+
 if __name__ == "__main__":
     uvicorn.run(app, host="localhost", port=8080)

poetry.lock

@@ -29,6 +29,7 @@ optional = true
 [tool.poetry.group.dev.dependencies]
 pyright = "^1.1.302"
 pytest = "^7.3.0"
+pytest-asyncio = "^0.21.0"
 black = "^23.3.0"
 
 [tool.poetry.group.examples]

pyproject.toml

@@ -1,11 +1,13 @@
 """Autometrics module."""
 import time
+import inspect
+
 from functools import wraps
-from typing import overload, TypeVar, Callable, Optional
+from typing import overload, TypeVar, Callable, Optional, Awaitable
 from typing_extensions import ParamSpec
 from .objectives import Objective
 from .tracker import get_tracker, Result
-from .utils import get_module_name, get_caller_function, write_docs
+from .utils import get_module_name, get_caller_function, append_docs_to_docstring
 
 
 P = ParamSpec("P")
@@ -18,6 +20,11 @@ def autometrics(func: Callable[P, T]) -> Callable[P, T]:
     ...
 
 
+@overload
+def autometrics(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+    ...
+
+
 # Decorator with arguments
 @overload
 def autometrics(*, objective: Optional[Objective] = None) -> Callable:
@@ -29,49 +36,100 @@ def autometrics(
     *,
     objective: Optional[Objective] = None,
 ):
-    """Decorator for tracking function calls and duration."""
+    """Decorator for tracking function calls and duration. Supports synchronous and async functions."""
+
+    def track_result_ok(start_time: float, function: str, module: str, caller: str):
+        get_tracker().finish(
+            start_time,
+            function=function,
+            module=module,
+            caller=caller,
+            objective=objective,
+            result=Result.OK,
+        )
+
+    def track_result_error(
+        start_time: float,
+        function: str,
+        module: str,
+        caller: str,
+    ):
+        get_tracker().finish(
+            start_time,
+            function=function,
+            module=module,
+            caller=caller,
+            objective=objective,
+            result=Result.ERROR,
+        )
+
+    def sync_decorator(func: Callable[P, T]) -> Callable[P, T]:
+        """Helper for decorating synchronous functions, to track calls and duration."""
 
-    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         module_name = get_module_name(func)
         func_name = func.__name__
 
         @wraps(func)
-        def wrapper(*args: P.args, **kwds: P.kwargs) -> T:
+        def sync_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
             start_time = time.time()
             caller = get_caller_function()
 
             try:
                 result = func(*args, **kwds)
-                get_tracker().finish(
+                track_result_ok(
+                    start_time, function=func_name, module=module_name, caller=caller
+                )
+
+            except Exception as exception:
+                result = exception.__class__.__name__
+                track_result_error(
                     start_time,
                     function=func_name,
                     module=module_name,
                     caller=caller,
-                    objective=objective,
-                    result=Result.OK,
+                )
+                # Reraise exception
+                raise exception
+            return result
+
+        sync_wrapper.__doc__ = append_docs_to_docstring(func, func_name, module_name)
+        return sync_wrapper
+
+    def async_decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+        """Helper for decorating async functions, to track calls and duration."""
+
+        module_name = get_module_name(func)
+        func_name = func.__name__
+
+        @wraps(func)
+        async def async_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
+            start_time = time.time()
+            caller = get_caller_function()
+
+            try:
+                result = await func(*args, **kwds)
+                track_result_ok(
+                    start_time, function=func_name, module=module_name, caller=caller
                 )
 
             except Exception as exception:
                 result = exception.__class__.__name__
-                get_tracker().finish(
+                track_result_error(
                     start_time,
                     function=func_name,
                     module=module_name,
                     caller=caller,
-                    objective=objective,
-                    result=Result.ERROR,
                 )
                 # Reraise exception
                 raise exception
             return result
 
-        if func.__doc__ is None:
-            wrapper.__doc__ = write_docs(func_name, module_name)
-        else:
-            wrapper.__doc__ = f"{func.__doc__}\n{write_docs(func_name, module_name)}"
-        return wrapper
+        async_wrapper.__doc__ = append_docs_to_docstring(func, func_name, module_name)
+        return async_wrapper
 
     if func is None:
-        return decorator
+        return sync_decorator
+    elif inspect.iscoroutinefunction(func):
+        return async_decorator(func)
     else:
-        return decorator(func)
+        return sync_decorator(func)

src/autometrics/decorator.py

@@ -1,5 +1,6 @@
 """Test the autometrics decorator."""
 import time
+import asyncio
 from prometheus_client.exposition import generate_latest
 import pytest
 
@@ -21,6 +22,18 @@ def error_function():
     raise RuntimeError("This is a test error")
 
 
+async def basic_async_function(sleep_duration: float = 1.0):
+    """This is a basic async function."""
+    await asyncio.sleep(sleep_duration)
+    return True
+
+
+async def error_async_function():
+    """This is an async function that raises an error."""
+    await asyncio.sleep(0.5)
+    raise RuntimeError("This is a test error")
+
+
 tracker_types = [TrackerType.PROMETHEUS, TrackerType.OPENTELEMETRY]
 
 
@@ -62,6 +75,39 @@ def test_basic(self):
         duration_sum = f"""function_calls_duration_sum{{function="{function_name}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
         assert duration_sum in data
 
+    @pytest.mark.asyncio
+    async def test_basic_async(self):
+        """This is a basic test."""
+
+        # set up the function + basic variables
+        caller = get_caller_function(depth=1)
+        assert caller is not None
+        assert caller != ""
+        function_name = basic_async_function.__name__
+        wrapped_function = autometrics(basic_async_function)
+
+        # Test that the function is *still* async after we wrap it
+        assert asyncio.iscoroutinefunction(wrapped_function) == True
+
+        await wrapped_function()
+
+        blob = generate_latest()
+        assert blob is not None
+        data = blob.decode("utf-8")
+
+        total_count = f"""function_calls_count_total{{caller="{caller}",function="{function_name}",module="test_decorator",objective_name="",objective_percentile="",result="ok"}} 1.0"""
+        assert total_count in data
+
+        for latency in ObjectiveLatency:
+            query = f"""function_calls_duration_bucket{{function="{function_name}",le="{latency.value}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
+            assert query in data
+
+        duration_count = f"""function_calls_duration_count{{function="{function_name}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
+        assert duration_count in data
+
+        duration_sum = f"""function_calls_duration_sum{{function="{function_name}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
+        assert duration_sum in data
+
     def test_objectives(self):
         """This is a test that covers objectives."""
 
@@ -132,3 +178,38 @@ def test_exception(self):
 
         duration_sum = f"""function_calls_duration_sum{{function="{function_name}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
         assert duration_sum in data
+
+    @pytest.mark.asyncio
+    async def test_async_exception(self):
+        """This is a test that covers exceptions."""
+        caller = get_caller_function(depth=1)
+        assert caller is not None
+        assert caller != ""
+
+        function_name = error_async_function.__name__
+        wrapped_function = autometrics(error_async_function)
+
+        # Test that the function is *still* async after we wrap it
+        assert asyncio.iscoroutinefunction(wrapped_function) == True
+
+        with pytest.raises(RuntimeError) as exception:
+            await wrapped_function()
+        assert "This is a test error" in str(exception.value)
+
+        # get the metrics
+        blob = generate_latest()
+        assert blob is not None
+        data = blob.decode("utf-8")
+
+        total_count = f"""function_calls_count_total{{caller="{caller}",function="{function_name}",module="test_decorator",objective_name="",objective_percentile="",result="error"}} 1.0"""
+        assert total_count in data
+
+        for latency in ObjectiveLatency:
+            query = f"""function_calls_duration_bucket{{function="{function_name}",le="{latency.value}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
+            assert query in data
+
+        duration_count = f"""function_calls_duration_count{{function="{function_name}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
+        assert duration_count in data
+
+        duration_sum = f"""function_calls_duration_sum{{function="{function_name}",module="test_decorator",objective_latency_threshold="",objective_name="",objective_percentile=""}}"""
+        assert duration_sum in data