Merge pull request #155 from codeflash-ai/context-import-bug

Context import bug

Author: misrasaurabh1

code_to_optimize/code_directories/retriever/import_test.py

@@ -0,0 +1,5 @@
+
+import code_to_optimize.code_directories.retriever.main
+
+def function_to_optimize():
+    return code_to_optimize.code_directories.retriever.main.fetch_and_transform_data()

codeflash/context/code_context_extractor.py

@@ -14,6 +14,7 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects
 from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages
+from codeflash.context.unused_definition_remover import remove_unused_definitions_by_function_names
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import (
     CodeContextType,
@@ -182,14 +183,16 @@ def extract_code_string_context_from_files(
             helpers_of_helpers_qualified_names = {
                 func.qualified_name for func in helpers_of_helpers.get(file_path, set())
             }
+            code_without_unused_defs = remove_unused_definitions_by_function_names(
+                original_code, qualified_function_names | helpers_of_helpers_qualified_names
+            )
             code_context = parse_code_and_prune_cst(
-                original_code,
+                code_without_unused_defs,
                 code_context_type,
                 qualified_function_names,
                 helpers_of_helpers_qualified_names,
                 remove_docstrings,
             )
-
         except ValueError as e:
             logger.debug(f"Error while getting read-only code: {e}")
             continue
@@ -214,8 +217,9 @@ def extract_code_string_context_from_files(
             continue
         try:
             qualified_helper_function_names = {func.qualified_name for func in helper_function_sources}
+            code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, qualified_helper_function_names)
             code_context = parse_code_and_prune_cst(
-                original_code, code_context_type, set(), qualified_helper_function_names, remove_docstrings
+                code_without_unused_defs, code_context_type, set(), qualified_helper_function_names, remove_docstrings
             )
         except ValueError as e:
             logger.debug(f"Error while getting read-only code: {e}")
@@ -283,8 +287,11 @@ def extract_code_markdown_context_from_files(
             helpers_of_helpers_qualified_names = {
                 func.qualified_name for func in helpers_of_helpers.get(file_path, set())
             }
+            code_without_unused_defs = remove_unused_definitions_by_function_names(
+                original_code, qualified_function_names | helpers_of_helpers_qualified_names
+            )
             code_context = parse_code_and_prune_cst(
-                original_code,
+                code_without_unused_defs,
                 code_context_type,
                 qualified_function_names,
                 helpers_of_helpers_qualified_names,
@@ -318,8 +325,9 @@ def extract_code_markdown_context_from_files(
             continue
         try:
             qualified_helper_function_names = {func.qualified_name for func in helper_function_sources}
+            code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, qualified_helper_function_names)
             code_context = parse_code_and_prune_cst(
-                original_code, code_context_type, set(), qualified_helper_function_names, remove_docstrings
+                code_without_unused_defs, code_context_type, set(), qualified_helper_function_names, remove_docstrings
             )
         except ValueError as e:
             logger.debug(f"Error while getting read-only code: {e}")