Warn on Using Span Inside Generator Functions

logfire/_internal/main.py

@@ -186,6 +186,7 @@ def _span(
         _span_name: str | None = None,
         _level: LevelName | int | None = None,
         _links: Sequence[tuple[SpanContext, otel_types.Attributes]] = (),
+        _warn_if_inside_generator: bool = True,
     ) -> LogfireSpan:
         try:
             if _level is not None:
@@ -196,6 +197,35 @@ def _span(
             else:
                 level_attributes = None
 
+            # we go two levels back to find the caller frame, as this method is called by logfire.span() method
+            caller_frame = inspect.currentframe().f_back.f_back  # type: ignore
+
+            # check if the caller is a generator function by checking the co_flags attribute of the code object
+            # and doing bit-wise AND checking for CO_GENERATOR or CO_ASYNC_GENERATOR value match
+            caller_is_generator = bool(
+                caller_frame and caller_frame.f_code.co_flags & (inspect.CO_GENERATOR | inspect.CO_ASYNC_GENERATOR)
+            )
+
+            is_from_context_manager = False
+
+            # Check if this call is coming from inside a context manager generator by inspecting call stack frames
+            if caller_is_generator:
+                previous_frame, origin_frame = inspect.currentframe().f_back, caller_frame.f_back  # type: ignore
+
+                for frame in [previous_frame, caller_frame, origin_frame]:
+                    code_name = frame.f_code.co_name  # type: ignore
+                    if code_name in ['__enter__', '__aenter__']:
+                        is_from_context_manager = True
+                        break
+
+            # usage within the context manager lifespan is legitimate, since the context manager controls the span's
+            # lifespan, and will ensure proper separation of concerns
+            if caller_is_generator and _warn_if_inside_generator and not is_from_context_manager:
+                warnings.warn(
+                    'Span is inside a generator function. See https://logfire.pydantic.dev/docs/reference/advanced/generators/#move-the-span-outside-the-generator.',
+                    RuntimeWarning,
+                )
+
             stack_info = get_user_stack_info()
             merged_attributes = {**stack_info, **attributes}
 
@@ -538,6 +568,7 @@ def span(
         _span_name: str | None = None,
         _level: LevelName | None = None,
         _links: Sequence[tuple[SpanContext, otel_types.Attributes]] = (),
+        _warn_if_inside_generator: bool = True,
         **attributes: Any,
     ) -> LogfireSpan:
         """Context manager for creating a span.
@@ -557,18 +588,21 @@ def span(
             _tags: An optional sequence of tags to include in the span.
             _level: An optional log level name.
             _links: An optional sequence of links to other spans. Each link is a tuple of a span context and attributes.
+            _warn_if_inside_generator: Set to `False` to prevent a warning when instrumenting a generator function.
             attributes: The arguments to include in the span and format the message template with.
                 Attributes starting with an underscore are not allowed.
         """
         if any(k.startswith('_') for k in attributes):
             raise ValueError('Attribute keys cannot start with an underscore.')
+
         return self._span(
             msg_template,
             attributes,
             _tags=_tags,
             _span_name=_span_name,
             _level=_level,
             _links=_links,
+            _warn_if_inside_generator=_warn_if_inside_generator,
         )
 
     @overload

tests/test_logfire.py

@@ -3543,3 +3543,71 @@ def test_warn_if_not_initialized_category():
 
         assert warnings_list[0].category == LogfireNotConfiguredWarning
         assert issubclass(LogfireNotConfiguredWarning, UserWarning)
+
+
+def test_warn_if_span_inside_generator():
+    """Test that warning is issued when a span is created inside a generator."""
+
+    def generator():
+        with logfire.span('span inside generator'):
+            yield
+
+    with pytest.warns(RuntimeWarning) as warnings_list:
+        next(generator())
+
+        assert warnings_list[0].category is RuntimeWarning
+        assert (
+            str(warnings_list[0].message)
+            == 'Span is inside a generator function. See https://logfire.pydantic.dev/docs/reference/advanced/generators/#move-the-span-outside-the-generator.'
+        )
+
+
+def test_no_warn_if_span_inside_generator():
+    """Test that warning is not issued when a span is created inside a generator with the
+    _warn_if_inside_generator option disabled."""
+
+    def generator():
+        with logfire.span('span inside generator', _warn_if_inside_generator=False):
+            yield
+
+    with warnings.catch_warnings(record=True) as warnings_list:
+        warnings.simplefilter('always')
+        next(generator())
+
+        assert len(warnings_list) == 0
+
+
+@pytest.mark.anyio
+async def test_warn_if_span_inside_async_generator():
+    """Test that warning is issued when a span is created inside an async generator."""
+
+    async def async_generator():
+        with logfire.span('span inside async generator'):
+            yield
+
+    with pytest.warns(RuntimeWarning) as warnings_list:
+        # we can replace this with global anext() when 3.9 is deprecated
+        await async_generator().__anext__()
+
+        assert warnings_list[0].category is RuntimeWarning
+        assert (
+            str(warnings_list[0].message)
+            == 'Span is inside a generator function. See https://logfire.pydantic.dev/docs/reference/advanced/generators/#move-the-span-outside-the-generator.'
+        )
+
+
+@pytest.mark.anyio
+async def test_no_warn_if_span_inside_async_generator():
+    """Test that warning is not issued when a span is created inside an async generator with the
+    _warn_if_inside_generator option disabled."""
+
+    async def async_generator():
+        with logfire.span('span inside async generator', _warn_if_inside_generator=False):
+            yield
+
+    with warnings.catch_warnings(record=True) as warnings_list:
+        warnings.simplefilter('always')
+        # we can replace this with global anext() when 3.9 is deprecated
+        await async_generator().__anext__()
+
+        assert len(warnings_list) == 0