# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
import functools
import asyncio
from typing import Any, Callable, Optional, Dict
try:
# pylint: disable = no-name-in-module
from opentelemetry import trace as opentelemetry_trace
tracer = opentelemetry_trace.get_tracer(__name__)
_tracing_library_available = True
except ModuleNotFoundError:
_tracing_library_available = False
if _tracing_library_available:
def trace_function(span_name: Optional[str] = None):
"""
A decorator for tracing function calls using OpenTelemetry.
This decorator handles various data types for function parameters and return values,
and records them as attributes in the trace span. The supported data types include:
- Basic data types: str, int, float, bool
- Collections: list, dict, tuple, set
Special handling for collections:
- If a collection (list, dict, tuple, set) contains nested collections, the entire collection
is converted to a string before being recorded as an attribute.
- Sets and dictionaries are always converted to strings to ensure compatibility with span attributes.
Object types are omitted, and the corresponding parameter is not traced.
:param span_name: The name of the span. If not provided, the function name is used.
:type span_name: Optional[str]
:return: The decorated function with tracing enabled.
:rtype: Callable
"""
def decorator(func: Callable) -> Callable:
@functools.wraps(func)
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
"""
Wrapper function for asynchronous functions.
:param args: Positional arguments passed to the function.
:type args: Tuple[Any]
:return: The result of the decorated asynchronous function.
:rtype: Any
"""
name = span_name if span_name else func.__name__
with tracer.start_as_current_span(name) as span:
try:
# Sanitize parameters and set them as attributes
sanitized_params = sanitize_parameters(func, *args, **kwargs)
span.set_attributes(sanitized_params)
result = await func(*args, **kwargs)
sanitized_result = sanitize_for_attributes(result)
if sanitized_result is not None:
if isinstance(sanitized_result, (list, dict, tuple, set)):
if any(isinstance(i, (list, dict, tuple, set)) for i in sanitized_result):
sanitized_result = str(sanitized_result)
span.set_attribute("code.function.return.value", sanitized_result) # type: ignore
return result
except Exception as e:
span.record_exception(e)
span.set_attribute("error.type", e.__class__.__qualname__) # type: ignore
raise
@functools.wraps(func)
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
"""
Wrapper function for synchronous functions.
:param args: Positional arguments passed to the function.
:type args: Tuple[Any]
:return: The result of the decorated synchronous function.
:rtype: Any
"""
name = span_name if span_name else func.__name__
with tracer.start_as_current_span(name) as span:
try:
# Sanitize parameters and set them as attributes
sanitized_params = sanitize_parameters(func, *args, **kwargs)
span.set_attributes(sanitized_params)
result = func(*args, **kwargs)
sanitized_result = sanitize_for_attributes(result)
if sanitized_result is not None:
if isinstance(sanitized_result, (list, dict, tuple, set)):
if any(isinstance(i, (list, dict, tuple, set)) for i in sanitized_result):
sanitized_result = str(sanitized_result)
span.set_attribute("code.function.return.value", sanitized_result) # type: ignore
return result
except Exception as e:
span.record_exception(e)
span.set_attribute("error.type", e.__class__.__qualname__) # type: ignore
raise
# Determine if the function is async
if asyncio.iscoroutinefunction(func):
return async_wrapper
return sync_wrapper
return decorator
else:
# Define a no-op decorator if OpenTelemetry is not available
[docs]
def trace_function(span_name: Optional[str] = None): # pylint: disable=unused-argument
"""
A no-op decorator for tracing function calls when OpenTelemetry is not available.
:param span_name: Not used in this version.
:type span_name: Optional[str]
:return: The original function.
:rtype: Callable
"""
def decorator(func: Callable) -> Callable:
return func
return decorator
def sanitize_parameters(func, *args, **kwargs) -> Dict[str, Any]:
"""
Sanitize function parameters to include only built-in data types.
:param func: The function being decorated.
:type func: Callable
:param args: Positional arguments passed to the function.
:type args: Tuple[Any]
:return: A dictionary of sanitized parameters.
:rtype: Dict[str, Any]
"""
import inspect
params = inspect.signature(func).parameters
sanitized_params = {}
for i, (name, param) in enumerate(params.items()):
if param.default == inspect.Parameter.empty and i < len(args):
value = args[i]
else:
value = kwargs.get(name, param.default)
sanitized_value = sanitize_for_attributes(value)
# Check if the collection has nested collections
if isinstance(sanitized_value, (list, dict, tuple, set)):
if any(isinstance(i, (list, dict, tuple, set)) for i in sanitized_value):
sanitized_value = str(sanitized_value)
if sanitized_value is not None:
sanitized_params["code.function.parameter." + name] = sanitized_value
return sanitized_params
# pylint: disable=R0911
def sanitize_for_attributes(value: Any, is_recursive: bool = False) -> Any:
"""
Sanitize a value to be used as an attribute.
:param value: The value to sanitize.
:type value: Any
:param is_recursive: Indicates if the function is being called recursively. Default is False.
:type is_recursive: bool
:return: The sanitized value or None if the value is not a supported type.
:rtype: Any
"""
if isinstance(value, (str, int, float, bool)):
return value
if isinstance(value, list):
return [
sanitize_for_attributes(item, True)
for item in value
if isinstance(item, (str, int, float, bool, list, dict, tuple, set))
]
if isinstance(value, dict):
retval = {
k: sanitize_for_attributes(v, True)
for k, v in value.items()
if isinstance(v, (str, int, float, bool, list, dict, tuple, set))
}
# dict to compatible with span attribute, so return it as a string
if is_recursive:
return retval
return str(retval)
if isinstance(value, tuple):
return tuple(
sanitize_for_attributes(item, True)
for item in value
if isinstance(item, (str, int, float, bool, list, dict, tuple, set))
)
if isinstance(value, set):
retval_set = {
sanitize_for_attributes(item, True)
for item in value
if isinstance(item, (str, int, float, bool, list, dict, tuple, set))
}
if is_recursive:
return retval_set
return str(retval_set)
return None