From d109be0b8ecd07a4cfbf292cbb1a96fbd3e7148e Mon Sep 17 00:00:00 2001 From: Dhruv Ahuja <83733638+dhruv-ahuja@users.noreply.github.com> Date: Thu, 25 Sep 2025 20:27:40 +0530 Subject: [PATCH 1/4] chore: warn if span()'s caller function is a generator --- logfire/_internal/main.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 4bf17a544..91faf917f 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,20 @@ 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) + ) + + if caller_is_generator and _warn_if_inside_generator: + 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 +553,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 +573,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 +587,7 @@ def span( _span_name=_span_name, _level=_level, _links=_links, + _warn_if_inside_generator=_warn_if_inside_generator, ) @overload From 16488d2d515f3478306b2ef571e1d6a7f986573d Mon Sep 17 00:00:00 2001 From: Dhruv Ahuja <83733638+dhruv-ahuja@users.noreply.github.com> Date: Thu, 25 Sep 2025 21:05:47 +0530 Subject: [PATCH 2/4] feat: add tests --- tests/test_logfire.py | 68 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) 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 From d5b5095132d7854ea62dd82b6e5e6e9b8a643fe2 Mon Sep 17 00:00:00 2001 From: Dhruv Ahuja <83733638+dhruv-ahuja@users.noreply.github.com> Date: Sat, 4 Oct 2025 23:19:35 +0530 Subject: [PATCH 3/4] fix: ensure warning only on non-context manager usage --- logfire/_internal/main.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 91faf917f..56718ecae 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -199,13 +199,32 @@ def _span( # 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) ) - if caller_is_generator and _warn_if_inside_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 = inspect.currentframe().f_back # type: ignore + origin_frame = caller_frame.f_back if caller_frame else None + + for frame in [previous_frame, caller_frame, origin_frame]: + if not frame: + continue + + code_name = frame.f_code.co_name + 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, From d57fc44d37c3b15e4697c57ab11758ba0c07b993 Mon Sep 17 00:00:00 2001 From: Dhruv Ahuja <83733638+dhruv-ahuja@users.noreply.github.com> Date: Sat, 4 Oct 2025 23:27:47 +0530 Subject: [PATCH 4/4] fix: remove unneeded code to prevent coverage error --- logfire/_internal/main.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 56718ecae..693edead4 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -210,14 +210,10 @@ def _span( # Check if this call is coming from inside a context manager generator by inspecting call stack frames if caller_is_generator: - previous_frame = inspect.currentframe().f_back # type: ignore - origin_frame = caller_frame.f_back if caller_frame else None + previous_frame, origin_frame = inspect.currentframe().f_back, caller_frame.f_back # type: ignore for frame in [previous_frame, caller_frame, origin_frame]: - if not frame: - continue - - code_name = frame.f_code.co_name + code_name = frame.f_code.co_name # type: ignore if code_name in ['__enter__', '__aenter__']: is_from_context_manager = True break