feat(tracer): add support for generators (#13377)

# PR Description

**Fix support for wrapping generator functions with `tracer.wrap()` and
ensure they generate traces in Datadog**
Related issue:
[#5403](#5403)

---

## Summary of the problem

Currently, when a generator function is wrapped with `tracer.wrap()`,
two major issues arise:

1. **Accessing the current span fails:**  
Inside the wrapped generator, `tracer.current_span()` returns `None`.
As a result, any attempt to set tags or interact with the span raises:

   ```
   AttributeError: 'NoneType' object has no attribute 'set_tag'
   ```

   Example:
   ```python
   @tracer.wrap()
   def foobar():
       current_span = tracer.current_span()
       current_span.set_tag("hello", "world")
       yield 1
   ```

   Calling `list(foobar())` → crashes with `AttributeError`.

2. **Traces are reported to Datadog with incorrect duration:**  
Even if the generator runs without explicit span interaction, traces
emitted to Datadog do not correctly report the execution time of the
function.
This is because the `tracer.wrap()` decorator does not maintain the span
context during generator iteration (`next()` or `async for`), so the
span gets opened and closed at the same time.

---

## Root cause

- The `wrap()` logic does not correctly handle Python generators (`def
... yield`) or async generators (`async def ... yield`).
- Without this, both local span interactions (`current_span()`) and
backend reporting (sending traces to Datadog) break.

---

## Proposed fix

This PR updates the `tracer.wrap()` decorator to:

- Add proper handling for **sync generators**:
  - Ensure `tracer.current_span()` is available.
  - Finalize the span after the generator is exhausted or on error.
  - Report traces to Datadog as expected.

- Add dedicated support for **async generators**:
  - Use an `async for` wrapper.
  - Maintain the tracing context.

With this change:
- Spans inside generators work (`current_span()` is valid).
- Traces from generator functions are correctly sent to Datadog.

---

## How to reproduce + verify the fix

Minimal reproducible example:
```python
@tracer.wrap()
def foobar():
    current_span = tracer.current_span()
    current_span.set_tag("hello", "world")
    yield 1

assert list(foobar()) == [1]
```

**Expected result:**
- No errors.
- Span is created and tagged.
- Trace is visible in Datadog with the tag `hello: world`.

**Added test cases:**
- Sync generator with span tags.
- Async generator with span tags.
- Backward compatibility for sync/async functions.

---

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: Brett Langdon <[email protected]>

Author: chasleslr

ddtrace/_trace/tracer.py

@@ -1,5 +1,6 @@
 from contextlib import contextmanager
 import functools
+import inspect
 from inspect import iscoroutinefunction
 from itertools import chain
 import logging
@@ -775,6 +776,56 @@ def flush(self):
         """Flush the buffer of the trace writer. This does nothing if an unbuffered trace writer is used."""
         self._span_aggregator.writer.flush_queue()
 
+    def _wrap_generator(
+        self,
+        f: AnyCallable,
+        span_name: str,
+        service: Optional[str] = None,
+        resource: Optional[str] = None,
+        span_type: Optional[str] = None,
+    ) -> AnyCallable:
+        """Wrap a generator function with tracing."""
+
+        @functools.wraps(f)
+        def func_wrapper(*args, **kwargs):
+            if getattr(self, "_wrap_executor", None):
+                return self._wrap_executor(
+                    self,
+                    f,
+                    args,
+                    kwargs,
+                    span_name,
+                    service=service,
+                    resource=resource,
+                    span_type=span_type,
+                )
+
+            with self.trace(span_name, service=service, resource=resource, span_type=span_type):
+                gen = f(*args, **kwargs)
+                for value in gen:
+                    yield value
+
+        return func_wrapper
+
+    def _wrap_generator_async(
+        self,
+        f: AnyCallable,
+        span_name: str,
+        service: Optional[str] = None,
+        resource: Optional[str] = None,
+        span_type: Optional[str] = None,
+    ) -> AnyCallable:
+        """Wrap a generator function with tracing."""
+
+        @functools.wraps(f)
+        async def func_wrapper(*args, **kwargs):
+            with self.trace(span_name, service=service, resource=resource, span_type=span_type):
+                agen = f(*args, **kwargs)
+                async for value in agen:
+                    yield value
+
+        return func_wrapper
+
     def wrap(
         self,
         name: Optional[str] = None,
@@ -812,6 +863,15 @@ async def coroutine():
             def coroutine():
                 return 'executed'
 
+        >>> # or use it on generators
+            @tracer.wrap()
+            def gen():
+                yield 'executed'
+
+        >>> @tracer.wrap()
+            async def gen():
+                yield 'executed'
+
         You can access the current span using `tracer.current_span()` to set
         tags:
 
@@ -825,10 +885,26 @@ def wrap_decorator(f: AnyCallable) -> AnyCallable:
             # FIXME[matt] include the class name for methods.
             span_name = name if name else "%s.%s" % (f.__module__, f.__name__)
 
-            # detect if the the given function is a coroutine to use the
-            # right decorator; this initial check ensures that the
+            # detect if the the given function is a coroutine and/or a generator
+            # to use the right decorator; this initial check ensures that the
             # evaluation is done only once for each @tracer.wrap
-            if iscoroutinefunction(f):
+            if inspect.isgeneratorfunction(f):
+                func_wrapper = self._wrap_generator(
+                    f,
+                    span_name,
+                    service=service,
+                    resource=resource,
+                    span_type=span_type,
+                )
+            elif inspect.isasyncgenfunction(f):
+                func_wrapper = self._wrap_generator_async(
+                    f,
+                    span_name,
+                    service=service,
+                    resource=resource,
+                    span_type=span_type,
+                )
+            elif iscoroutinefunction(f):
                 # create an async wrapper that awaits the coroutine and traces it
                 @functools.wraps(f)
                 async def func_wrapper(*args, **kwargs):

releasenotes/notes/fix-tracing-generator-043422ae1d1974aa.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    tracing: Fixes support for wrapping generator and async generator functions with `tracer.wrap()`. Previously, calling `tracer.current_span()` inside a wrapped generator function would return `None`, leading to `AttributeError` when interacting with the span. Additionally, traces reported to Datadog showed incorrect durations, as span context was not maintained across generator iteration. This change ensures that `tracer.wrap()` now correctly handles both sync and async generators by preserving the tracing context throughout their execution and finalizing spans correctly. Users can now safely use `tracer.current_span()` within generator functions and expect accurate trace reporting.

tests/contrib/asyncio/test_tracer.py

@@ -1,4 +1,5 @@
 """Ensure that the tracer works with asynchronous executions within the same ``IOLoop``."""
+
 import asyncio
 import os
 import re
@@ -223,3 +224,31 @@ async def my_function():
     rb"created at .*/dd-trace-py/ddtrace/contrib/internal/asyncio/patch.py:.* took .* seconds"
     match = re.match(pattern, err)
     assert match, err
+
+
+@pytest.mark.asyncio
+async def test_wrapped_generator(tracer):
+    @tracer.wrap("decorated_generator", service="s", resource="r", span_type="t")
+    async def f(tag_name, tag_value):
+        # make sure we can still set tags
+        span = tracer.current_span()
+        span.set_tag(tag_name, tag_value)
+
+        for i in range(3):
+            yield i
+
+    result = [item async for item in f("a", "b")]
+    assert result == [0, 1, 2]
+
+    traces = tracer.pop_traces()
+
+    assert 1 == len(traces)
+    spans = traces[0]
+    assert 1 == len(spans)
+    span = spans[0]
+
+    assert span.name == "decorated_generator"
+    assert span.service == "s"
+    assert span.resource == "r"
+    assert span.span_type == "t"
+    assert span.get_tag("a") == "b"

tests/tracer/test_tracer.py

@@ -8,6 +8,7 @@
 import logging
 from os import getpid
 import threading
+import time
 from unittest.case import SkipTest
 
 import mock
@@ -284,6 +285,34 @@ def wrapped_function(param, kw_param=None):
             (dict(name="wrap.overwrite", service="webserver", meta=dict(args="(42,)", kwargs="{'kw_param': 42}")),),
         )
 
+    def test_tracer_wrap_generator(self):
+        @self.tracer.wrap("decorated_generator", service="s", resource="r", span_type="t")
+        def f(tag_name, tag_value):
+            # make sure we can still set tags
+            span = self.tracer.current_span()
+            span.set_tag(tag_name, tag_value)
+
+            for i in range(3):
+                time.sleep(0.01)
+                yield i
+
+        result = list(f("a", "b"))
+        assert result == [0, 1, 2]
+
+        self.assert_span_count(1)
+        span = self.get_root_span()
+        span.assert_matches(
+            name="decorated_generator",
+            service="s",
+            resource="r",
+            span_type="t",
+            meta=dict(a="b"),
+        )
+
+        # tracer should finish _after_ the generator has been exhausted
+        assert span.duration is not None
+        assert span.duration > 0.03
+
     def test_tracer_disabled(self):
         self.tracer.enabled = True
         with self.trace("foo") as s: