33import ast
44from collections import defaultdict
55from dataclasses import dataclass , field
6- from pathlib import Path
7- from typing import Optional
6+ from typing import TYPE_CHECKING
7+
8+ if TYPE_CHECKING :
9+ from pathlib import Path
10+ from typing import TYPE_CHECKING , Optional
811
912import libcst as cst
1013
1114from codeflash .cli_cmds .console import logger
1215from codeflash .code_utils .code_replacer import replace_function_definitions_in_module
13- from codeflash .models .models import CodeOptimizationContext , FunctionSource
16+
17+ if TYPE_CHECKING :
18+ from codeflash .discovery .functions_to_optimize import FunctionToOptimize
19+ from codeflash .models .models import CodeOptimizationContext , FunctionSource
1420
1521
1622@dataclass
@@ -493,11 +499,12 @@ def print_definitions(definitions: dict[str, UsageInfo]) -> None:
493499
494500
495501def revert_unused_helper_functions (
496- project_root , unused_helpers : list [FunctionSource ], original_helper_code : dict [Path , str ]
502+ project_root : Path , unused_helpers : list [FunctionSource ], original_helper_code : dict [Path , str ]
497503) -> None :
498504 """Revert unused helper functions back to their original definitions.
499505
500506 Args:
507+ project_root: project_root
501508 unused_helpers: List of unused helper functions to revert
502509 original_helper_code: Dictionary mapping file paths to their original code
503510
@@ -516,9 +523,6 @@ def revert_unused_helper_functions(
516523 for file_path , helpers_in_file in unused_helpers_by_file .items ():
517524 if file_path in original_helper_code :
518525 try :
519- # Read current file content
520- current_code = file_path .read_text (encoding = "utf8" )
521-
522526 # Get original code for this file
523527 original_code = original_helper_code [file_path ]
524528
@@ -557,7 +561,6 @@ def _analyze_imports_in_optimized_code(
557561 # Precompute a two-level dict: module_name -> func_name -> [helpers]
558562 helpers_by_file_and_func = defaultdict (dict )
559563 helpers_by_file = defaultdict (list ) # preserved for "import module"
560- helpers_append = helpers_by_file_and_func .setdefault
561564 for helper in code_context .helper_functions :
562565 jedi_type = helper .jedi_definition .type
563566 if jedi_type != "class" :
@@ -606,11 +609,12 @@ def _analyze_imports_in_optimized_code(
606609
607610
608611def detect_unused_helper_functions (
609- function_to_optimize , code_context : CodeOptimizationContext , optimized_code : str
612+ function_to_optimize : FunctionToOptimize , code_context : CodeOptimizationContext , optimized_code : str
610613) -> list [FunctionSource ]:
611614 """Detect helper functions that are no longer called by the optimized entrypoint function.
612615
613616 Args:
617+ function_to_optimize: The function to optimize
614618 code_context: The code optimization context containing helper functions
615619 optimized_code: The optimized code to analyze
616620
@@ -702,8 +706,9 @@ def detect_unused_helper_functions(
702706 logger .debug (f"Helper function { helper_qualified_name } is still called in optimized code" )
703707 logger .debug (f" Called via: { possible_call_names .intersection (called_function_names )} " )
704708
705- return unused_helpers
709+ ret_val = unused_helpers
706710
707711 except Exception as e :
708712 logger .debug (f"Error detecting unused helper functions: { e } " )
709- return []
713+ ret_val = []
714+ return ret_val
0 commit comments