1414from codeflash .cli_cmds .console import logger
1515from codeflash .code_utils .code_extractor import add_needed_imports_from_module , find_preexisting_objects
1616from codeflash .code_utils .code_utils import get_qualified_name , path_belongs_to_site_packages
17+ from codeflash .context .unused_definition_remover import remove_unused_definitions_by_function_names
1718from codeflash .discovery .functions_to_optimize import FunctionToOptimize
1819from codeflash .models .models import (
1920 CodeContextType ,
@@ -182,14 +183,16 @@ def extract_code_string_context_from_files(
182183 helpers_of_helpers_qualified_names = {
183184 func .qualified_name for func in helpers_of_helpers .get (file_path , set ())
184185 }
186+ code_without_unused_defs = remove_unused_definitions_by_function_names (
187+ original_code , qualified_function_names | helpers_of_helpers_qualified_names
188+ )
185189 code_context = parse_code_and_prune_cst (
186- original_code ,
190+ code_without_unused_defs ,
187191 code_context_type ,
188192 qualified_function_names ,
189193 helpers_of_helpers_qualified_names ,
190194 remove_docstrings ,
191195 )
192-
193196 except ValueError as e :
194197 logger .debug (f"Error while getting read-only code: { e } " )
195198 continue
@@ -214,8 +217,9 @@ def extract_code_string_context_from_files(
214217 continue
215218 try :
216219 qualified_helper_function_names = {func .qualified_name for func in helper_function_sources }
220+ code_without_unused_defs = remove_unused_definitions_by_function_names (original_code , qualified_helper_function_names )
217221 code_context = parse_code_and_prune_cst (
218- original_code , code_context_type , set (), qualified_helper_function_names , remove_docstrings
222+ code_without_unused_defs , code_context_type , set (), qualified_helper_function_names , remove_docstrings
219223 )
220224 except ValueError as e :
221225 logger .debug (f"Error while getting read-only code: { e } " )
@@ -283,8 +287,11 @@ def extract_code_markdown_context_from_files(
283287 helpers_of_helpers_qualified_names = {
284288 func .qualified_name for func in helpers_of_helpers .get (file_path , set ())
285289 }
290+ code_without_unused_defs = remove_unused_definitions_by_function_names (
291+ original_code , qualified_function_names | helpers_of_helpers_qualified_names
292+ )
286293 code_context = parse_code_and_prune_cst (
287- original_code ,
294+ code_without_unused_defs ,
288295 code_context_type ,
289296 qualified_function_names ,
290297 helpers_of_helpers_qualified_names ,
@@ -318,8 +325,9 @@ def extract_code_markdown_context_from_files(
318325 continue
319326 try :
320327 qualified_helper_function_names = {func .qualified_name for func in helper_function_sources }
328+ code_without_unused_defs = remove_unused_definitions_by_function_names (original_code , qualified_helper_function_names )
321329 code_context = parse_code_and_prune_cst (
322- original_code , code_context_type , set (), qualified_helper_function_names , remove_docstrings
330+ code_without_unused_defs , code_context_type , set (), qualified_helper_function_names , remove_docstrings
323331 )
324332 except ValueError as e :
325333 logger .debug (f"Error while getting read-only code: { e } " )
0 commit comments