Update code_extractor.py

Author: KRRT7

codeflash/code_utils/code_extractor.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
 import ast
-from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Optional, Set
+from typing import TYPE_CHECKING, Optional
 
 import libcst as cst
 import libcst.matchers as m
@@ -11,23 +10,25 @@
 from libcst.helpers import calculate_module_and_package
 
 from codeflash.cli_cmds.console import logger
-from codeflash.models.models import FunctionParent, FunctionSource
+from codeflash.models.models import FunctionParent
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from libcst.helpers import ModuleNameAndPackage
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+    from codeflash.models.models import FunctionSource
 
-from typing import List
 
 
 class GlobalAssignmentCollector(cst.CSTVisitor):
     """Collects all global assignment statements."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
-        self.assignments: Dict[str, cst.Assign] = {}
-        self.assignment_order: List[str] = []
+        self.assignments: dict[str, cst.Assign] = {}
+        self.assignment_order: list[str] = []
         # Track scope depth to identify global assignments
         self.scope_depth = 0
         self.if_else_depth = 0
@@ -72,11 +73,11 @@ def visit_Assign(self, node: cst.Assign) -> Optional[bool]:
 class GlobalAssignmentTransformer(cst.CSTTransformer):
     """Transforms global assignments in the original file with those from the new file."""
 
-    def __init__(self, new_assignments: Dict[str, cst.Assign], new_assignment_order: List[str]):
+    def __init__(self, new_assignments: dict[str, cst.Assign], new_assignment_order: list[str]) -> None:
         super().__init__()
         self.new_assignments = new_assignments
         self.new_assignment_order = new_assignment_order
-        self.processed_assignments: Set[str] = set()
+        self.processed_assignments: set[str] = set()
         self.scope_depth = 0
         self.if_else_depth = 0
 
@@ -145,7 +146,7 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
 class GlobalStatementCollector(cst.CSTVisitor):
     """Visitor that collects all global statements (excluding imports and functions/classes)."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.global_statements = []
         self.in_function_or_class = False
@@ -178,7 +179,7 @@ def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
 class LastImportFinder(cst.CSTVisitor):
     """Finds the position of the last import statement in the module."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.last_import_line = 0
         self.current_line = 0
@@ -193,7 +194,7 @@ def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
 class ImportInserter(cst.CSTTransformer):
     """Transformer that inserts global statements after the last import."""
 
-    def __init__(self, global_statements: List[cst.SimpleStatementLine], last_import_line: int):
+    def __init__(self, global_statements: list[cst.SimpleStatementLine], last_import_line: int) -> None:
         super().__init__()
         self.global_statements = global_statements
         self.last_import_line = last_import_line
@@ -208,7 +209,7 @@ def leave_SimpleStatementLine(
         # If we're right after the last import and haven't inserted yet
         if self.current_line == self.last_import_line and not self.inserted:
             self.inserted = True
-            return cst.Module(body=[updated_node] + self.global_statements)
+            return cst.Module(body=[updated_node, *self.global_statements])
 
         return cst.Module(body=[updated_node])
 
@@ -222,7 +223,7 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         return updated_node
 
 
-def extract_global_statements(source_code: str) -> List[cst.SimpleStatementLine]:
+def extract_global_statements(source_code: str) -> list[cst.SimpleStatementLine]:
     """Extract global statements from source code."""
     module = cst.parse_module(source_code)
     collector = GlobalStatementCollector()
@@ -285,8 +286,7 @@ def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
     transformer = GlobalAssignmentTransformer(new_collector.assignments, new_collector.assignment_order)
     transformed_module = original_module.visit(transformer)
 
-    dst_module_code = transformed_module.code
-    return dst_module_code
+    return transformed_module.code
 
 
 def add_needed_imports_from_module(