Don't auto-trace generators (#386)

Author: alexmojaki

docs/guides/advanced/generators.md

@@ -90,6 +90,8 @@ Traceback (most recent call last):
 ValueError: <Token var=<ContextVar name='current_context' default={} at 0x10afa3f60> at 0x10de034c0> was created in a different Context
 ```
 
+This is why generator functions are not traced by [`logfire.install_auto_tracing()`][logfire.Logfire.install_auto_tracing].
+
 ## What you can do
 
 ### Move the span outside the generator

docs/guides/onboarding_checklist/add_auto_tracing.md

@@ -1,6 +1,6 @@
 # Auto-tracing
 
-The [`logfire.install_auto_tracing`][logfire.Logfire.install_auto_tracing]
+The [`logfire.install_auto_tracing()`][logfire.Logfire.install_auto_tracing] method
 will trace all function calls in the specified modules.
 
 This works by changing how those modules are imported,
@@ -21,6 +21,9 @@ from app.main import main
 main()
 ```
 
+!!! note
+    Generator functions will not be traced for reasons explained [here](../advanced/generators.md).
+
 ## Filtering modules to trace
 
 The `modules` argument can be a list of module names.

logfire/_internal/auto_trace/rewrite_ast.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import ast
+import inspect
+import types
 import uuid
 from dataclasses import dataclass
 from functools import partial
@@ -94,6 +96,12 @@ def visit_FunctionDef(self, node: ast.FunctionDef | ast.AsyncFunctionDef):
 
         return super().visit_FunctionDef(node)
 
+    def rewrite_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef, qualname: str) -> ast.AST:
+        if is_generator_function(node):
+            return node
+
+        return super().rewrite_function(node, qualname)
+
     def logfire_method_call_node(self, node: ast.FunctionDef | ast.AsyncFunctionDef, qualname: str) -> ast.Call:
         # See the exec_source docstring
         index = len(self.context_factories)
@@ -158,3 +166,15 @@ def no_auto_trace(x: T) -> T:
     This decorator simply returns the argument unchanged, so there is zero runtime overhead.
     """
     return x  # pragma: no cover
+
+
+GENERATOR_CODE_FLAGS = inspect.CO_GENERATOR | inspect.CO_ASYNC_GENERATOR
+
+
+def is_generator_function(func_def: ast.FunctionDef | ast.AsyncFunctionDef):
+    module_node = ast.parse('')
+    module_node.body = [func_def]
+    code = compile(module_node, '<string>', 'exec')
+    return any(
+        isinstance(const, types.CodeType) and (const.co_flags & GENERATOR_CODE_FLAGS) for const in code.co_consts
+    )

logfire/_internal/main.py

@@ -780,12 +780,14 @@ def install_auto_tracing(
     ) -> None:
         """Install automatic tracing.
 
-        This will trace all function calls in the modules specified by the modules argument.
+        This will trace all non-generator function calls in the modules specified by the modules argument.
         It's equivalent to wrapping the body of every function in matching modules in `with logfire.span(...):`.
 
         !!! note
             This function MUST be called before any of the modules to be traced are imported.
 
+            Generator functions will not be traced for reasons explained [here](https://docs.pydantic.dev/logfire/guides/advanced/generators/).
+
         This works by inserting a new meta path finder into `sys.meta_path`, so inserting another finder before it
         may prevent it from working.
 

tests/test_auto_trace.py

@@ -85,45 +85,12 @@ def test_auto_trace_sample(exporter: TestExporter) -> None:
                     'logfire.pending_parent_id': '0000000000000001',
                 },
             },
-            {
-                'name': 'Calling tests.auto_trace_samples.foo.gen (pending)',
-                'context': {'trace_id': 1, 'span_id': 6, 'is_remote': False},
-                'parent': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
-                'start_time': 3000000000,
-                'end_time': 3000000000,
-                'attributes': {
-                    'code.filepath': 'foo.py',
-                    'code.lineno': 123,
-                    'code.function': 'gen',
-                    'logfire.msg_template': 'Calling tests.auto_trace_samples.foo.gen',
-                    'logfire.msg': 'Calling tests.auto_trace_samples.foo.gen',
-                    'logfire.span_type': 'pending_span',
-                    'logfire.tags': ('auto-trace',),
-                    'logfire.pending_parent_id': '0000000000000003',
-                },
-            },
-            {
-                'name': 'Calling tests.auto_trace_samples.foo.gen',
-                'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
-                'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
-                'start_time': 3000000000,
-                'end_time': 4000000000,
-                'attributes': {
-                    'code.filepath': 'foo.py',
-                    'code.lineno': 123,
-                    'code.function': 'gen',
-                    'logfire.msg_template': 'Calling tests.auto_trace_samples.foo.gen',
-                    'logfire.span_type': 'span',
-                    'logfire.tags': ('auto-trace',),
-                    'logfire.msg': 'Calling tests.auto_trace_samples.foo.gen',
-                },
-            },
             {
                 'name': 'Calling async_gen via @instrument',
                 'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
                 'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                 'start_time': 2000000000,
-                'end_time': 5000000000,
+                'end_time': 3000000000,
                 'attributes': {
                     'code.filepath': 'foo.py',
                     'code.lineno': 123,
@@ -138,7 +105,7 @@ def test_auto_trace_sample(exporter: TestExporter) -> None:
                 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                 'parent': None,
                 'start_time': 1000000000,
-                'end_time': 7000000000,
+                'end_time': 5000000000,
                 'attributes': {
                     'code.filepath': 'foo.py',
                     'code.lineno': 123,
@@ -152,7 +119,7 @@ def test_auto_trace_sample(exporter: TestExporter) -> None:
                 'events': [
                     {
                         'name': 'exception',
-                        'timestamp': 6000000000,
+                        'timestamp': 4000000000,
                         'attributes': {
                             'exception.type': 'IndexError',
                             'exception.message': 'list index out of range',
@@ -443,6 +410,26 @@ def test_no_auto_trace():
     )
 
 
+# language=Python
+generators_sample = """
+def make_gen():
+    def gen():
+        async def foo():
+            async def bar():
+                pass
+            yield bar()
+        yield from foo()
+    return gen
+"""
+
+
+def test_generators():
+    assert get_calling_strings(generators_sample) == {
+        'Calling module.name.make_gen',
+        'Calling module.name.make_gen.<locals>.gen.<locals>.foo.<locals>.bar',
+    }
+
+
 def test_min_duration(exporter: TestExporter):
     install_auto_tracing('tests.auto_trace_samples.simple_nesting', min_duration=5)