Skip to content

Commit d3ae2c7

Browse files
committed
Attempt 2 to reduce cognitive complexity
1 parent 2d4d670 commit d3ae2c7

File tree

1 file changed

+39
-35
lines changed

1 file changed

+39
-35
lines changed

src/common/tracing.py

Lines changed: 39 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
import asyncio
22
from functools import wraps
3-
43
from opentelemetry import trace
54

6-
# Get the _tracer instance (You can set your own _tracer name)
75
tracer = trace.get_tracer(__name__)
86

97

@@ -16,41 +14,47 @@ def trace_function(trace_attributes: bool = True, trace_result: bool = True):
1614
- trace_result (bool): If False, disables adding the function's result to the span.
1715
"""
1816

17+
def set_span_attributes(span, args, kwargs, result=None):
18+
"""Helper to set function arguments and results as span attributes."""
19+
if trace_attributes:
20+
span.set_attribute("function.args", str(args))
21+
span.set_attribute("function.kwargs", str(kwargs))
22+
if trace_result and result is not None:
23+
span.set_attribute("function.result", str(result))
24+
25+
def record_exception(span, exception):
26+
"""Helper to handle exception recording in a span."""
27+
span.record_exception(exception)
28+
span.set_status(trace.status.Status(trace.status.StatusCode.ERROR))
29+
30+
async def handle_async(span, func, *args, **kwargs):
31+
"""Handle asynchronous functions."""
32+
try:
33+
result = await func(*args, **kwargs)
34+
set_span_attributes(span, args, kwargs, result)
35+
return result
36+
except Exception as e:
37+
record_exception(span, e)
38+
raise
39+
40+
def handle_sync(span, func, *args, **kwargs):
41+
"""Handle synchronous functions."""
42+
try:
43+
result = func(*args, **kwargs)
44+
set_span_attributes(span, args, kwargs, result)
45+
return result
46+
except Exception as e:
47+
record_exception(span, e)
48+
raise
49+
1950
def decorator(func):
2051
@wraps(func)
21-
def sync_or_async_wrapper(*args, **kwargs):
52+
def wrapper(*args, **kwargs):
2253
with tracer.start_as_current_span(func.__name__) as span:
23-
try:
24-
# Set function arguments as attributes
25-
if trace_attributes:
26-
span.set_attribute("function.args", str(args))
27-
span.set_attribute("function.kwargs", str(kwargs))
28-
29-
async def async_handler():
30-
result = await func(*args, **kwargs)
31-
# Add result to span
32-
if trace_result:
33-
span.set_attribute("function.result", str(result))
34-
return result
35-
36-
def sync_handler():
37-
result = func(*args, **kwargs)
38-
# Add result to span
39-
if trace_result:
40-
span.set_attribute("function.result", str(result))
41-
return result
42-
43-
if asyncio.iscoroutinefunction(func):
44-
return async_handler()
45-
else:
46-
return sync_handler()
47-
48-
except Exception as e:
49-
# Record the exception in the span
50-
span.record_exception(e)
51-
span.set_status(trace.status.Status(trace.status.StatusCode.ERROR))
52-
raise
53-
54-
return sync_or_async_wrapper
54+
if asyncio.iscoroutinefunction(func):
55+
return handle_async(span, func, *args, **kwargs)
56+
return handle_sync(span, func, *args, **kwargs)
57+
58+
return wrapper
5559

5660
return decorator

0 commit comments

Comments
 (0)