Fix error in checking for generators in auto-tracing (#498)

logfire-api/logfire_api/_internal/auto_trace/rewrite_ast.pyi

@@ -1,7 +1,6 @@
 import ast
 from ..ast_utils import BaseTransformer as BaseTransformer, LogfireArgs as LogfireArgs
 from ..main import Logfire as Logfire
-from _typeshed import Incomplete
 from dataclasses import dataclass
 from typing import Any, Callable, ContextManager, TypeVar
 
@@ -51,7 +50,4 @@ def no_auto_trace(x: T) -> T:
 
     This decorator simply returns the argument unchanged, so there is zero runtime overhead.
     """
-
-GENERATOR_CODE_FLAGS: Incomplete
-
-def is_generator_function(func_def: ast.FunctionDef | ast.AsyncFunctionDef): ...
+def has_yield(node: ast.AST): ...

logfire/_internal/auto_trace/rewrite_ast.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
 import ast
-import inspect
-import types
 import uuid
+from collections import deque
 from dataclasses import dataclass
 from functools import partial
 from typing import TYPE_CHECKING, Any, Callable, ContextManager, TypeVar
@@ -97,7 +96,7 @@ 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):
+        if has_yield(node):
             return node
 
         return super().rewrite_function(node, qualname)
@@ -168,13 +167,12 @@ def no_auto_trace(x: T) -> T:
     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
-    )
+def has_yield(node: ast.AST):
+    queue = deque([node])
+    while queue:
+        node = queue.popleft()
+        for child in ast.iter_child_nodes(node):
+            if isinstance(child, (ast.Yield, ast.YieldFrom)):
+                return True
+            if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
+                queue.append(child)

tests/test_auto_trace.py

@@ -417,13 +417,12 @@ def test_no_auto_trace():
     )
 
 
-# language=Python
 generators_sample = """
 def make_gen():
     def gen():
         async def foo():
             async def bar():
-                pass
+                return lambda: (yield 1)
             yield bar()
         yield from foo()
     return gen