Improve caller detection

src/autometrics/decorator.py

@@ -1,13 +1,18 @@
 """Autometrics module."""
 import time
 import inspect
-
 from functools import wraps
 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, append_docs_to_docstring
+from .utils import (
+    get_function_name,
+    get_module_name,
+    get_caller_function,
+    append_docs_to_docstring,
+)
 
 
 P = ParamSpec("P")
@@ -85,13 +90,13 @@ def sync_decorator(func: Callable[P, T]) -> Callable[P, T]:
         """Helper for decorating synchronous functions, to track calls and duration."""
 
         module_name = get_module_name(func)
-        func_name = func.__name__
+        func_name = get_function_name(func)
         register_function_info(func_name, module_name)
 
         @wraps(func)
         def sync_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
             start_time = time.time()
-            caller = get_caller_function()
+            caller = get_caller_function(func_name=func_name, module_name=module_name)
 
             try:
                 if track_concurrency:
@@ -120,13 +125,13 @@ 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__
+        func_name = get_function_name(func)
         register_function_info(func_name, module_name)
 
         @wraps(func)
         async def async_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
             start_time = time.time()
-            caller = get_caller_function()
+            caller = get_caller_function(func_name=func_name, module_name=module_name)
 
             try:
                 if track_concurrency:

src/autometrics/test_caller.py

@@ -0,0 +1,63 @@
+"""Tests for caller tracking."""
+from functools import wraps
+from prometheus_client.exposition import generate_latest
+
+from .decorator import autometrics
+from .utils import get_caller_function, get_function_name, get_module_name
+
+
+def test_caller_detection():
+    """This is a test to see if the caller is properly detected."""
+
+    def dummy_decorator(func):
+        @wraps(func)
+        def dummy_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        return dummy_wrapper
+
+    def another_decorator(func):
+        @wraps(func)
+        def another_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        return another_wrapper
+
+    @dummy_decorator
+    @autometrics
+    @another_decorator
+    def foo():
+        print(get_caller_function(func_name="foo", module_name=__name__))
+        pass
+
+    foo()
+
+    func_name = get_function_name(foo)
+
+    blob = generate_latest()
+    assert blob is not None
+    data = blob.decode("utf-8")
+
+    expected = """function_calls_count_total{caller="test_caller_detection",function="test_caller_detection.<locals>.foo",module="autometrics.test_caller",objective_name="",objective_percentile="",result="ok"} 1.0"""
+    assert "wrapper" not in data
+    assert expected in data
+
+
+def test_caller_detection_external():
+    """Same as above, but for external wrapper."""
+
+    from .test_caller_external import bar
+
+    bar()
+
+    caller = get_caller_function(depth=1)
+    func_name = get_function_name(bar)
+    module_name = get_module_name(bar)
+
+    blob = generate_latest()
+    assert blob is not None
+    data = blob.decode("utf-8")
+
+    expected = f"""function_calls_count_total{{caller="{caller}",function="{func_name}",module="{module_name}",objective_name="",objective_percentile="",result="ok"}} 1.0"""
+    assert "external_wrapper" not in data
+    assert expected in data

src/autometrics/test_caller_external.py

@@ -0,0 +1,27 @@
+"""Supplementary for caller tests."""
+from functools import wraps
+
+from .decorator import autometrics
+
+
+def dummy_external_decorator(func):
+    @wraps(func)
+    def dummy_external_wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return dummy_external_wrapper
+
+
+def another_external_decorator(func):
+    @wraps(func)
+    def another_external_wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return another_external_wrapper
+
+
+@dummy_external_decorator
+@autometrics
+@another_external_decorator
+def bar():
+    pass