diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 4bf17a544..693edead4 100644 --- a/logfire/_internal/main.py +++ b/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,11 +588,13 @@ 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, @@ -569,6 +602,7 @@ def span( _span_name=_span_name, _level=_level, _links=_links, + _warn_if_inside_generator=_warn_if_inside_generator, ) @overload diff --git a/tests/test_logfire.py b/tests/test_logfire.py index dd53ef44d..36caef2bc 100644 --- a/tests/test_logfire.py +++ b/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