diff --git a/ddtrace/_trace/tracer.py b/ddtrace/_trace/tracer.py index 2a68a7d49c2..647cfc59b74 100644 --- a/ddtrace/_trace/tracer.py +++ b/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 diff --git a/releasenotes/notes/fix-tracer-wrap-generator-return-value-fbaec09a9670e06b.yaml b/releasenotes/notes/fix-tracer-wrap-generator-return-value-fbaec09a9670e06b.yaml new file mode 100644 index 00000000000..ada2d9a1d27 --- /dev/null +++ b/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. \ No newline at end of file diff --git a/tests/tracer/test_tracer.py b/tests/tracer/test_tracer.py index f6f58bcac20..77935ecada6 100644 --- a/tests/tracer/test_tracer.py +++ b/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: