diff --git a/logfire/_internal/auto_trace/import_hook.py b/logfire/_internal/auto_trace/import_hook.py index 5603b8122..6579a9452 100644 --- a/logfire/_internal/auto_trace/import_hook.py +++ b/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. diff --git a/logfire/_internal/auto_trace/rewrite_ast.py b/logfire/_internal/auto_trace/rewrite_ast.py index 986c641de..1ec71d194 100644 --- a/logfire/_internal/auto_trace/rewrite_ast.py +++ b/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_`. @@ -34,16 +36,20 @@ 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, @@ -51,7 +57,6 @@ def rewrite_ast( 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 diff --git a/logfire/_internal/stack_info.py b/logfire/_internal/stack_info.py index a19904060..9205c9076 100644 --- a/logfire/_internal/stack_info.py +++ b/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 != '': + if code.co_name != '': # 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,6 +89,7 @@ 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. @@ -96,6 +97,7 @@ def is_user_code(code: CodeType) -> bool: """ return not ( str(Path(code.co_filename).absolute()).startswith(NON_USER_CODE_PREFIXES) + or code.co_filename.startswith('<') or code.co_name in ('', '', '') ) diff --git a/tests/test_auto_trace.py b/tests/test_auto_trace.py index 18dbbff04..11b194c84 100644 --- a/tests/test_auto_trace.py +++ b/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', diff --git a/tests/test_source_code_extraction.py b/tests/test_source_code_extraction.py index 77dc7ea26..407e389e7 100644 --- a/tests/test_source_code_extraction.py +++ b/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': '', - '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': '', - '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': '', - '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',