Skip to content

Commit 886616f

Browse files
fix unused helper tests
1 parent 330bf91 commit 886616f

File tree

3 files changed

+43
-23
lines changed

3 files changed

+43
-23
lines changed

codeflash/models/models.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ def markdown(self) -> str:
170170
)
171171

172172
def path_to_code_string(self) -> dict[str, str]:
173-
return {code_string.file_path: code_string.code for code_string in self.code_strings}
173+
return {str(code_string.file_path): code_string.code for code_string in self.code_strings}
174174

175175
@staticmethod
176176
def from_str_with_markers(code_with_markers: str) -> CodeStringsMarkdown:

codeflash/optimization/function_optimizer.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -627,16 +627,19 @@ def replace_function_and_helpers_with_optimized_code(
627627
if helper_function.jedi_definition.type != "class":
628628
read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name)
629629
for module_abspath, qualified_names in read_writable_functions_by_file_path.items():
630-
relative_module_path = module_abspath.relative_to(self.project_root)
630+
relative_module_path = str(module_abspath.relative_to(self.project_root))
631631
logger.debug(f"applying optimized code to: {relative_module_path}")
632632

633-
optimized_code = file_to_code_context.get(relative_module_path)
634-
if not optimized_code:
635-
msg = f"Optimized code not found for {relative_module_path}, existing files in the context are: {list(file_to_code_context.keys())}, re-check your 'split markers'"
636-
raise ValueError(msg)
633+
scoped_optimized_code = file_to_code_context.get(relative_module_path, None)
634+
if scoped_optimized_code is None:
635+
logger.warning(
636+
f"Optimized code not found for {relative_module_path}, existing files in the context are: {list(file_to_code_context.keys())}, re-check your 'split markers'"
637+
)
638+
scoped_optimized_code = ""
639+
637640
did_update |= replace_function_definitions_in_module(
638641
function_names=list(qualified_names),
639-
optimized_code=optimized_code,
642+
optimized_code=scoped_optimized_code,
640643
module_abspath=module_abspath,
641644
preexisting_objects=code_context.preexisting_objects,
642645
project_root_path=self.project_root,

tests/test_unused_helper_revert.py

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import pytest
77
from codeflash.context.unused_definition_remover import detect_unused_helper_functions
88
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
9+
from codeflash.models.models import get_code_block_splitter
910
from codeflash.optimization.function_optimizer import FunctionOptimizer
1011
from codeflash.verification.verification_utils import TestConfig
1112

@@ -55,7 +56,8 @@ def test_detect_unused_helper_functions(temp_project):
5556
temp_dir, main_file, test_cfg = temp_project
5657

5758
# Optimized version that only calls one helper
58-
optimized_code = """
59+
optimized_code = f"""
60+
{get_code_block_splitter("main.py")}
5961
def entrypoint_function(n):
6062
\"\"\"Optimized function that only calls one helper.\"\"\"
6163
result1 = helper_function_1(n)
@@ -99,7 +101,8 @@ def helper_function_2(x):
99101

100102
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
101103
# First modify the optimized code to include a MODIFIED unused helper
102-
optimized_code_with_modified_helper = """
104+
optimized_code_with_modified_helper = f"""
105+
{get_code_block_splitter("main.py")}
103106
def entrypoint_function(n):
104107
\"\"\"Optimized function that only calls one helper.\"\"\"
105108
result1 = helper_function_1(n)
@@ -159,7 +162,8 @@ def test_revert_unused_helper_functions(temp_project):
159162
temp_dir, main_file, test_cfg = temp_project
160163

161164
# Optimized version that only calls one helper and modifies the unused one
162-
optimized_code = """
165+
optimized_code = f"""
166+
{get_code_block_splitter("main.py") }
163167
def entrypoint_function(n):
164168
\"\"\"Optimized function that only calls one helper.\"\"\"
165169
result1 = helper_function_1(n)
@@ -221,7 +225,8 @@ def test_no_unused_helpers_no_revert(temp_project):
221225
temp_dir, main_file, test_cfg = temp_project
222226

223227
# Optimized version that still calls both helpers
224-
optimized_code = """
228+
optimized_code = f"""
229+
{get_code_block_splitter("main.py")}
225230
def entrypoint_function(n):
226231
\"\"\"Optimized function that still calls both helpers.\"\"\"
227232
result1 = helper_function_1(n)
@@ -303,13 +308,16 @@ def helper_function_2(x):
303308
""")
304309

305310
# Optimized version that only calls one helper
306-
optimized_code = """
311+
optimized_code = f"""
312+
{get_code_block_splitter("main.py")}
307313
from helpers import helper_function_1
308314
309315
def entrypoint_function(n):
310316
\"\"\"Optimized function that only calls one helper.\"\"\"
311317
result1 = helper_function_1(n)
312318
return result1 + n * 3 # Inlined helper_function_2
319+
320+
{get_code_block_splitter("helpers.py")}
313321
"""
314322

315323
# Create test config
@@ -384,7 +392,6 @@ def helper_function_2(x):
384392

385393
# Apply optimization and test reversion
386394
optimizer.replace_function_and_helpers_with_optimized_code(code_context, optimized_code, original_helper_code)
387-
388395
# Check main file content
389396
main_content = main_file.read_text()
390397
assert "result1 + n * 3" in main_content, "Entrypoint function should be optimized"
@@ -478,7 +485,8 @@ def helper_method_2(self, x):
478485
""")
479486

480487
# Optimized version that only calls one helper method
481-
optimized_code = """
488+
optimized_code = f"""
489+
{get_code_block_splitter("main.py") }
482490
class Calculator:
483491
def entrypoint_method(self, n):
484492
\"\"\"Optimized method that only calls one helper.\"\"\"
@@ -537,7 +545,8 @@ def helper_method_2(self, x):
537545

538546
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
539547
# Update optimized code to include a MODIFIED unused helper
540-
optimized_code_with_modified_helper = """
548+
optimized_code_with_modified_helper = f"""
549+
{get_code_block_splitter("main.py")}
541550
class Calculator:
542551
def entrypoint_method(self, n):
543552
\"\"\"Optimized method that only calls one helper.\"\"\"
@@ -619,7 +628,8 @@ def process_data(self, n):
619628
""")
620629

621630
# Optimized version that only calls one external helper
622-
optimized_code = """
631+
optimized_code = f"""
632+
{get_code_block_splitter("main.py") }
623633
def external_helper_1(x):
624634
\"\"\"External helper function.\"\"\"
625635
return x * 2
@@ -678,7 +688,8 @@ def process_data(self, n):
678688

679689
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
680690
# Update optimized code to include a MODIFIED unused helper
681-
optimized_code_with_modified_helper = """
691+
optimized_code_with_modified_helper = f"""
692+
{get_code_block_splitter("main.py")}
682693
def external_helper_1(x):
683694
\"\"\"External helper function.\"\"\"
684695
return x * 2
@@ -716,7 +727,8 @@ def process_data(self, n):
716727

717728
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
718729
# Update optimized code to include a MODIFIED unused helper
719-
optimized_code_with_modified_helper = """
730+
optimized_code_with_modified_helper = f"""
731+
{get_code_block_splitter("main.py")}
720732
def external_helper_1(x):
721733
\"\"\"External helper function.\"\"\"
722734
return x * 2
@@ -786,7 +798,8 @@ def local_helper(self, x):
786798
""")
787799

788800
# Optimized version that inlines one helper
789-
optimized_code = """
801+
optimized_code = f"""
802+
{get_code_block_splitter("main.py") }
790803
def global_helper_1(x):
791804
return x * 2
792805
@@ -954,7 +967,8 @@ def clean_data(x):
954967
""")
955968

956969
# Optimized version that only uses some functions
957-
optimized_code = """
970+
optimized_code = f"""
971+
{get_code_block_splitter("main.py") }
958972
import utils
959973
from math_helpers import add
960974
@@ -1115,7 +1129,8 @@ def divide_numbers(x, y):
11151129
""")
11161130

11171131
# Optimized version that only uses add_numbers
1118-
optimized_code = """
1132+
optimized_code = f"""
1133+
{get_code_block_splitter("main.py") }
11191134
import calculator
11201135
11211136
def entrypoint_function(n):
@@ -1317,7 +1332,8 @@ def calculate_class(cls, n):
13171332
""")
13181333

13191334
# Optimized static method that inlines one utility
1320-
optimized_static_code = """
1335+
optimized_static_code = f"""
1336+
{get_code_block_splitter("main.py")}
13211337
def utility_function_1(x):
13221338
return x * 2
13231339
@@ -1384,7 +1400,8 @@ def calculate_class(cls, n):
13841400

13851401
# Also test the complete replace_function_and_helpers_with_optimized_code workflow
13861402
# Update optimized code to include a MODIFIED unused helper
1387-
optimized_static_code_with_modified_helper = """
1403+
optimized_static_code_with_modified_helper = f"""
1404+
{get_code_block_splitter("main.py")}
13881405
def utility_function_1(x):
13891406
return x * 2
13901407

0 commit comments

Comments
 (0)