From 4a5057fadcde6e3775b199d7061003854e9e10f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kacper=20Kali=C5=84ski?=
 <47140412+KaQuMiQ@users.noreply.github.com>
Date: Mon, 28 Oct 2024 10:55:18 +0100
Subject: [PATCH] Add function tracing

---
 src/haiway/__init__.py         |  14 +++-
 src/haiway/helpers/__init__.py |   4 ++
 src/haiway/helpers/tracing.py  | 122 +++++++++++++++++++++++++++++++++
 3 files changed, 139 insertions(+), 1 deletion(-)
 create mode 100644 src/haiway/helpers/tracing.py

diff --git a/src/haiway/__init__.py b/src/haiway/__init__.py
index ad10729..56b7f3c 100644
--- a/src/haiway/__init__.py
+++ b/src/haiway/__init__.py
@@ -6,7 +6,16 @@
     ScopeMetrics,
     ctx,
 )
-from haiway.helpers import asynchronous, cache, retry, throttle, timeout
+from haiway.helpers import (
+    ArgumentsTrace,
+    ResultTrace,
+    asynchronous,
+    cache,
+    retry,
+    throttle,
+    timeout,
+    traced,
+)
 from haiway.state import State
 from haiway.types import (
     MISSING,
@@ -34,6 +43,7 @@
 
 __all__ = [
     "always",
+    "ArgumentsTrace",
     "async_always",
     "async_noop",
     "asynchronous",
@@ -57,11 +67,13 @@
     "MissingState",
     "noop",
     "not_missing",
+    "ResultTrace",
     "retry",
     "ScopeMetrics",
     "setup_logging",
     "State",
     "throttle",
     "timeout",
+    "traced",
     "when_missing",
 ]
diff --git a/src/haiway/helpers/__init__.py b/src/haiway/helpers/__init__.py
index d55ce39..2caf455 100644
--- a/src/haiway/helpers/__init__.py
+++ b/src/haiway/helpers/__init__.py
@@ -3,11 +3,15 @@
 from haiway.helpers.retries import retry
 from haiway.helpers.throttling import throttle
 from haiway.helpers.timeouted import timeout
+from haiway.helpers.tracing import ArgumentsTrace, ResultTrace, traced
 
 __all__ = [
+    "ArgumentsTrace",
     "asynchronous",
     "cache",
+    "ResultTrace",
     "retry",
     "throttle",
     "timeout",
+    "traced",
 ]
diff --git a/src/haiway/helpers/tracing.py b/src/haiway/helpers/tracing.py
new file mode 100644
index 0000000..e0ff913
--- /dev/null
+++ b/src/haiway/helpers/tracing.py
@@ -0,0 +1,122 @@
+from asyncio import iscoroutinefunction
+from collections.abc import Callable, Coroutine
+from typing import Any, Self, cast
+
+from haiway.context import ctx
+from haiway.state import State
+from haiway.types import MISSING, Missing
+from haiway.utils import mimic_function
+
+__all__ = [
+    "traced",
+    "ArgumentsTrace",
+    "ResultTrace",
+]
+
+
+class ArgumentsTrace(State):
+    if __debug__:
+
+        @classmethod
+        def of(cls, *args: Any, **kwargs: Any) -> Self:
+            return cls(
+                args=args if args else MISSING,
+                kwargs=kwargs if kwargs else MISSING,
+            )
+
+    else:  # remove tracing for non debug runs to prevent accidental secret leaks
+
+        @classmethod
+        def of(cls, *args: Any, **kwargs: Any) -> Self:
+            return cls(
+                args=MISSING,
+                kwargs=MISSING,
+            )
+
+    args: tuple[Any, ...] | Missing
+    kwargs: dict[str, Any] | Missing
+
+
+class ResultTrace(State):
+    if __debug__:
+
+        @classmethod
+        def of(
+            cls,
+            value: Any,
+            /,
+        ) -> Self:
+            return cls(result=value)
+
+    else:  # remove tracing for non debug runs to prevent accidental secret leaks
+
+        @classmethod
+        def of(
+            cls,
+            value: Any,
+            /,
+        ) -> Self:
+            return cls(result=MISSING)
+
+    result: Any | Missing
+
+
+def traced[**Args, Result](
+    function: Callable[Args, Result],
+    /,
+) -> Callable[Args, Result]:
+    if __debug__:
+        if iscoroutinefunction(function):
+            return cast(
+                Callable[Args, Result],
+                _traced_async(
+                    function,
+                    label=function.__name__,
+                ),
+            )
+        else:
+            return _traced_sync(
+                function,
+                label=function.__name__,
+            )
+
+    else:  # do not trace on non debug runs
+        return function
+
+
+def _traced_sync[**Args, Result](
+    function: Callable[Args, Result],
+    /,
+    label: str,
+) -> Callable[Args, Result]:
+    @mimic_function(function)
+    def wrapped(
+        *args: Args.args,
+        **kwargs: Args.kwargs,
+    ) -> Result:
+        with ctx.scope(label):
+            ctx.record(ArgumentsTrace.of(*args, **kwargs))
+            result: Result = function(*args, **kwargs)
+            ctx.record(ResultTrace.of(result))
+            return result
+
+    return wrapped
+
+
+def _traced_async[**Args, Result](
+    function: Callable[Args, Coroutine[Any, Any, Result]],
+    /,
+    label: str,
+) -> Callable[Args, Coroutine[Any, Any, Result]]:
+    @mimic_function(function)
+    async def wrapped(
+        *args: Args.args,
+        **kwargs: Args.kwargs,
+    ) -> Result:
+        with ctx.scope(label):
+            ctx.record(ArgumentsTrace.of(*args, **kwargs))
+            result: Result = await function(*args, **kwargs)
+            ctx.record(ResultTrace.of(result))
+            return result
+
+    return wrapped