11import ast
22from ..ast_utils import BaseTransformer as BaseTransformer , LogfireArgs as LogfireArgs
33from ..main import Logfire as Logfire
4+ from contextlib import AbstractContextManager as AbstractContextManager
45from dataclasses import dataclass
5- from typing import Any , Callable , ContextManager , TypeVar
6+ from typing import Any , Callable , TypeVar
67
78def compile_source (tree : ast .AST , filename : str , module_name : str , logfire_instance : Logfire , min_duration : int ) -> Callable [[dict [str , Any ]], None ]:
89 """Compile a modified AST of the module's source code in the module's namespace.
@@ -21,13 +22,13 @@ def compile_source(tree: ast.AST, filename: str, module_name: str, logfire_insta
2122 If `min_duration` is greater than 0, then `context_factories[index]` is initially `MeasureTime`.
2223 Otherwise, it's initially the `partial` above.
2324 """
24- def rewrite_ast (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 : ...
25+ def rewrite_ast (tree : ast .AST , filename : str , logfire_name : str , module_name : str , logfire_instance : Logfire , context_factories : list [Callable [[], AbstractContextManager [Any ]]], min_duration : int ) -> ast .AST : ...
2526
2627@dataclass
2728class AutoTraceTransformer (BaseTransformer ):
2829 """Trace all encountered functions except those explicitly marked with `@no_auto_trace`."""
2930 logfire_instance : Logfire
30- context_factories : list [Callable [[], ContextManager [Any ]]]
31+ context_factories : list [Callable [[], AbstractContextManager [Any ]]]
3132 min_duration : int
3233 def check_no_auto_trace (self , node : ast .FunctionDef | ast .AsyncFunctionDef | ast .ClassDef ) -> bool :
3334 """Return true if the node has a `@no_auto_trace` or `@logfire.no_auto_trace` decorator."""
0 commit comments