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