Handle errors in auto-tracing better

logfire/_internal/auto_trace/import_hook.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
+import ast
 import sys
 from dataclasses import dataclass
 from importlib.abc import Loader, MetaPathFinder
 from importlib.machinery import ModuleSpec
 from importlib.util import spec_from_loader
 from types import ModuleType
-from typing import TYPE_CHECKING, Callable, Iterator, Sequence, cast
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Sequence, cast
 
-from .rewrite_ast import exec_source
+from ..utils import log_internal_error
+from .rewrite_ast import compile_source
 from .types import AutoTraceModule
 
 if TYPE_CHECKING:
@@ -48,10 +50,38 @@ def find_spec(
             if not source:
                 continue
 
-            if not self.modules_filter(AutoTraceModule(fullname, plain_spec.origin)):
+            # We fully expect plain_spec.origin and self.get_filename(...)
+            # to be the same thing (a valid filename), but they're optional.
+            filename = plain_spec.origin
+            if not filename:  # pragma: no cover
+                try:
+                    filename = cast('str | None', plain_spec.loader.get_filename(fullname))  # type: ignore
+                except Exception:
+                    pass
+
+            if not self.modules_filter(AutoTraceModule(fullname, filename)):
                 return None  # tell the import system to try the next meta path finder
 
-            loader = LogfireLoader(plain_spec, source, self.logfire, self.min_duration)
+            try:
+                tree = ast.parse(source)
+            except Exception:  # pragma: no cover
+                # The plain finder gave us invalid source code. Try another one.
+                # A very likely case is that the source code really is invalid,
+                # in which case we'll eventually return None and the normal system will raise the error,
+                # giving the user a normal traceback instead of a confusing and ugly one mentioning logfire.
+                continue
+
+            filename = filename or f'<{fullname}>'
+
+            try:
+                execute = compile_source(tree, filename, fullname, self.logfire, self.min_duration)
+            except Exception:  # pragma: no cover
+                # Auto-tracing failed with an unexpected error. Ensure that this doesn't crash the whole application.
+                # This error handling is why we compile in the finder. Once we return a loader, we've committed to it.
+                log_internal_error()
+                return None  # tell the import system to try the next meta path finder
+
+            loader = LogfireLoader(plain_spec, execute)
             return spec_from_loader(fullname, loader)
 
     def _find_plain_specs(
@@ -79,29 +109,15 @@ class LogfireLoader(Loader):
     plain_spec: ModuleSpec
     """A spec for the module that was returned by another meta path finder (see `LogfireFinder._find_plain_specs`)."""
 
-    source: str
-    """The source code of the module, as returned by `plain_spec.loader.get_source(fullname)`."""
-
-    logfire: Logfire
-    min_duration: int
+    execute: Callable[[dict[str, Any]], None]
+    """A function which accepts module globals and executes the compiled code."""
 
     def exec_module(self, module: ModuleType):
         """Execute a modified AST of the module's source code in the module's namespace.
 
         This is called by the import system.
         """
-        # We fully expect self.plain_spec.origin, module.__file__, and self.get_filename(...)
-        # to all be the same thing (a valid filename), but technically they're all optional,
-        # so this is just an abundance of caution.
-        filename = self.plain_spec.origin or module.__file__
-        if not filename:  # pragma: no cover
-            try:
-                filename = self.get_filename(module.__name__)
-            except Exception:
-                pass
-        filename = filename or f'<{module.__name__}>'
-
-        exec_source(self.source, filename, module.__name__, module.__dict__, self.logfire, self.min_duration)
+        self.execute(module.__dict__)
 
     # This is required when `exec_module` is defined.
     # It returns None to indicate that the usual module creation process should be used.

logfire/_internal/auto_trace/rewrite_ast.py

@@ -15,10 +15,12 @@
     from ..main import Logfire
 
 
-def exec_source(
-    source: str, filename: str, module_name: str, globs: dict[str, Any], logfire_instance: Logfire, min_duration: int
-) -> None:
-    """Execute a modified AST of the module's source code in the module's namespace.
+def compile_source(
+    tree: ast.AST, filename: str, module_name: str, logfire_instance: Logfire, min_duration: int
+) -> Callable[[dict[str, Any]], None]:
+    """Compile a modified AST of the module's source code in the module's namespace.
+
+    Returns a function which accepts module globals and executes the compiled code.
 
     The modified AST wraps the body of every function definition in `with context_factories[index]():`.
     `context_factories` is added to the module's namespace as `logfire_<uuid>`.
@@ -34,24 +36,27 @@ def exec_source(
     """
     logfire_name = f'logfire_{uuid.uuid4().hex}'
     context_factories: list[Callable[[], ContextManager[Any]]] = []
-    globs[logfire_name] = context_factories
-    tree = rewrite_ast(source, filename, logfire_name, module_name, logfire_instance, context_factories, min_duration)
+    tree = rewrite_ast(tree, filename, logfire_name, module_name, logfire_instance, context_factories, min_duration)
     assert isinstance(tree, ast.Module)  # for type checking
     # dont_inherit=True is necessary to prevent the module from inheriting the __future__ import from this module.
     code = compile(tree, filename, 'exec', dont_inherit=True)
-    exec(code, globs, globs)
+
+    def execute(globs: dict[str, Any]):
+        globs[logfire_name] = context_factories
+        exec(code, globs, globs)
+
+    return execute
 
 
 def rewrite_ast(
-    source: str,
+    tree: ast.AST,
     filename: str,
     logfire_name: str,
     module_name: str,
     logfire_instance: Logfire,
     context_factories: list[Callable[[], ContextManager[Any]]],
     min_duration: int,
 ) -> ast.AST:
-    tree = ast.parse(source)
     logfire_args = LogfireArgs(logfire_instance._tags, logfire_instance._sample_rate)  # type: ignore
     transformer = AutoTraceTransformer(
         logfire_args, logfire_name, filename, module_name, logfire_instance, context_factories, min_duration

logfire/_internal/stack_info.py

@@ -39,7 +39,7 @@ def get_filepath_attribute(file: str) -> StackInfo:
 @lru_cache(maxsize=2048)
 def get_code_object_info(code: CodeType) -> StackInfo:
     result = get_filepath_attribute(code.co_filename)
-    if code.co_name != '<module>':
+    if code.co_name != '<module>':  # pragma: no branch
         result['code.function'] = code.co_qualname if sys.version_info >= (3, 11) else code.co_name
     result['code.lineno'] = code.co_firstlineno
     return result
@@ -89,13 +89,15 @@ def is_user_code(code: CodeType) -> bool:
         - the standard library
         - site-packages (specifically wherever opentelemetry is installed)
         - the logfire package
+        - an unknown location (e.g. a dynamically generated code object) indicated by a filename starting with '<'
     - It is a list/dict/set comprehension.
         These are artificial frames only created before Python 3.12,
         and they are always called directly from the enclosing function so it makes sense to skip them.
         On the other hand, generator expressions and lambdas might be called far away from where they are defined.
     """
     return not (
         str(Path(code.co_filename).absolute()).startswith(NON_USER_CODE_PREFIXES)
+        or code.co_filename.startswith('<')
         or code.co_name in ('<listcomp>', '<dictcomp>', '<setcomp>')
     )
 

tests/test_auto_trace.py

@@ -211,7 +211,7 @@ def only_ellipsis_function():
 def test_rewrite_ast():
     context_factories: list[Callable[[], ContextManager[Any]]] = []
     tree = rewrite_ast(
-        nested_sample,
+        ast.parse(nested_sample),
         'foo.py',
         'logfire_span',
         'module.name',
@@ -396,7 +396,7 @@ def definitely_not_traced_method(self):
 def get_calling_strings(sample: str):
     context_factories: list[Callable[[], ContextManager[Any]]] = []
     rewrite_ast(
-        sample,
+        ast.parse(sample),
         'foo.py',
         'logfire_span',
         'module.name',

tests/test_source_code_extraction.py

@@ -111,7 +111,7 @@ def test_source_code_extraction_module(exporter: TestExporter) -> None:
         )
 
     assert normalize_filepaths(
-        exporter.exported_spans_as_dict(strip_filepaths=False, fixed_line_number=None, _strip_function_qualname=False)
+        exporter.exported_spans_as_dict(strip_filepaths=False, _strip_function_qualname=False)
     ) == snapshot(
         [
             {
@@ -134,8 +134,9 @@ def test_source_code_extraction_module(exporter: TestExporter) -> None:
 Failed to introspect calling code. Please report this issue to Logfire. Falling back to normal message formatting which may result in loss of information if using an f-string. Set inspect_arguments=False in logfire.configure() to suppress this warning. The problem was:
 No source code available. This happens when running in an interactive shell, using exec(), or running .pyc files without the source .py files.\
 """,
-                    'code.filepath': '<string>',
-                    'code.lineno': 2,
+                    'code.filepath': 'tests/test_source_code_extraction.py',
+                    'code.function': 'test_source_code_extraction_module',
+                    'code.lineno': 123,
                 },
             },
             {
@@ -145,8 +146,9 @@ def test_source_code_extraction_module(exporter: TestExporter) -> None:
                 'start_time': 2000000000,
                 'end_time': 3000000000,
                 'attributes': {
-                    'code.filepath': '<string>',
-                    'code.lineno': 2,
+                    'code.filepath': 'tests/test_source_code_extraction.py',
+                    'code.function': 'test_source_code_extraction_module',
+                    'code.lineno': 123,
                     'logfire.msg_template': 'from module',
                     'logfire.msg': 'from module',
                     'logfire.span_type': 'span',
@@ -169,7 +171,7 @@ def test_source_code_extraction_exec_no_inspect_arguments(
     )
 
     assert normalize_filepaths(
-        exporter.exported_spans_as_dict(strip_filepaths=False, fixed_line_number=None, _strip_function_qualname=False)
+        exporter.exported_spans_as_dict(strip_filepaths=False, _strip_function_qualname=False)
     ) == snapshot(
         [
             {
@@ -179,8 +181,9 @@ def test_source_code_extraction_exec_no_inspect_arguments(
                 'start_time': 1000000000,
                 'end_time': 2000000000,
                 'attributes': {
-                    'code.filepath': '<string>',
-                    'code.lineno': 2,
+                    'code.filepath': 'tests/test_source_code_extraction.py',
+                    'code.function': 'test_source_code_extraction_exec_no_inspect_arguments',
+                    'code.lineno': 123,
                     'logfire.msg_template': 'from module',
                     'logfire.span_type': 'span',
                     'logfire.msg': 'from module',