fix(tracing): preserve generator return values in @tracer.wrap() decorator [backport 3.9]

Author: brettlangdon

ddtrace/_trace/tracer.py

@@ -775,7 +775,8 @@ def func_wrapper(*args, **kwargs):
                 )
 
             with self.trace(span_name, service=service, resource=resource, span_type=span_type):
-                yield from f(*args, **kwargs)
+                return_value = yield from f(*args, **kwargs)
+                return return_value
 
         return func_wrapper
 

releasenotes/notes/fix-tracer-wrap-generator-return-value-fbaec09a9670e06b.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    tracing: This fix resolves an issue where the ``@tracer.wrap()`` decorator failed to preserve return values from generator functions, causing ``StopIteration.value`` to be ``None`` instead of the actual returned value.

tests/tracer/test_tracer.py

@@ -355,6 +355,28 @@ def wrapper():
             "visible to the caller"
         )
 
+    def test_tracer_wrap_generator_with_return_value(self):
+        @self.tracer.wrap()
+        def iter_signals():
+            for i in range(10):
+                yield i
+            return 10
+
+        with self.trace("root") as span:
+            signals = iter_signals()
+            while True:
+                try:
+                    # DEV: We don't need the return value
+                    next(signals)
+                except StopIteration as e:
+                    assert e.value == 10
+                    span.set_metric("num_signals", e.value)
+                    break
+
+        self.assert_span_count(2)
+        root_span = self.get_root_span()
+        root_span.assert_matches(name="root", metrics={"num_signals": 10})
+
     def test_tracer_disabled(self):
         self.tracer.enabled = True
         with self.trace("foo") as s: