move _set_span_attribuets into span

Author: antonpirker

sentry_sdk/tracing.py

@@ -599,9 +599,24 @@ def set_tag(self, key, value):
         # type: (str, Any) -> None
         self._tags[key] = value
 
-    def set_data(self, key, value):
-        # type: (str, Any) -> None
-        self._data[key] = value
+    def set_data(self, key, value=None):
+        # type: (Union[str, Dict[str, Any]], Any) -> None
+        """Set data on the span.
+
+        Can be called in two ways:
+        - set_data(key, value) - sets a single key-value pair
+        - set_data({"key": "value"}) - sets multiple key-value pairs from a dict
+        """
+        if isinstance(key, dict):
+            # Dictionary calling pattern: set_data({"key": "value"})
+            for k, v in key.items():
+                self._data[k] = v
+        else:
+            # Traditional calling pattern: set_data(key, value)
+            if value is None:
+                raise ValueError("Value must be provided when key is a string")
+
+            self._data[key] = value
 
     def set_flag(self, flag, result):
         # type: (str, bool) -> None
@@ -1272,8 +1287,8 @@ def set_tag(self, key, value):
         # type: (str, Any) -> None
         pass
 
-    def set_data(self, key, value):
-        # type: (str, Any) -> None
+    def set_data(self, key, value=None):
+        # type: (Union[str, Dict[str, Any]], Any) -> None
         pass
 
     def set_status(self, value):
@@ -1343,18 +1358,20 @@ def trace(func):
         pass
 
 
-def trace(func=None, *, as_type="span", name=None, attributes=None):
-    # type: (Optional[Callable[P, R]], str, Optional[str], Optional[dict[str, Any]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
+def trace(func=None, *, template="span", op=None, name=None, attributes=None):
+    # type: (Optional[Callable[P, R]], str, Optional[str], Optional[str], Optional[dict[str, Any]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
     """
     Decorator to start a child span.
 
     :param func: The function to trace.
-    :param as_type: The type of span to create.
+    :param template: The type of span to create.
+    :param op: The operation of the span.
     :param name: The name of the span. (defaults to the function name)
     :param attributes: Additional attributes to set on the span.
     """
     decorator = create_span_decorator(
-        template=as_type,
+        template=template,
+        op=op,
         name=name,
         attributes=attributes,
     )

sentry_sdk/tracing_utils.py

@@ -864,20 +864,6 @@ def _get_span_op(template):
     return op
 
 
-def _set_span_attributes(span, attributes):
-    # type: (Any, Optional[dict[str, Any]]) -> None
-    """
-    Set the given attributes on the given span.
-
-    :param span: The span to set attributes on.
-    :param attributes: The attributes to set on the span.
-    """
-    attributes = attributes or {}
-
-    for key, value in attributes.items():
-        span.set_data(key, value)
-
-
 def _set_input_attributes(span, template, args, kwargs):
     # type: (Span, str, tuple[Any, ...], dict[str, Any]) -> None
     """
@@ -907,7 +893,7 @@ def _set_input_attributes(span, template, args, kwargs):
             SPANDATA.GEN_AI_TOOL_NAME: span.description,
         }
 
-    _set_span_attributes(span, attributes)
+    span.set_data(attributes)
 
 
 def _set_output_attributes(span, template, result):
@@ -923,7 +909,7 @@ def _set_output_attributes(span, template, result):
 
     # TODO: Implement :)
 
-    _set_span_attributes(span, attributes)
+    span.set_data(attributes)
 
 
 def create_span_decorator(template, op=None, name=None, attributes=None):
@@ -969,7 +955,7 @@ async def async_wrapper(*args, **kwargs):
                 result = await f(*args, **kwargs)
 
                 _set_output_attributes(span, template, result)
-                _set_span_attributes(span, attributes)
+                span.set_data(attributes)
 
                 return result
 
@@ -1004,7 +990,7 @@ def sync_wrapper(*args, **kwargs):
                 result = f(*args, **kwargs)
 
                 _set_output_attributes(span, template, result)
-                _set_span_attributes(span, attributes)
+                span.set_data(attributes)
 
                 return result