Merge branch 'main' into worktree/mirror-all-arg

Author: KRRT7

codeflash/code_utils/coverage_utils.py

@@ -12,11 +12,12 @@
 
 def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]:
     """Extract the single dependent function from the code context excluding the main function."""
-    ast_tree = ast.parse(code_context.testgen_context_code)
-
-    dependent_functions = {
-        node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
-    }
+    dependent_functions = set()
+    for code_string in code_context.testgen_context.code_strings:
+        ast_tree = ast.parse(code_string.code)
+        dependent_functions.update(
+            {node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))}
+        )
 
     if main_function in dependent_functions:
         dependent_functions.discard(main_function)

codeflash/context/code_context_extractor.py

@@ -114,32 +114,32 @@ def get_code_optimization_context(
             read_only_context_code = ""
 
     # Extract code context for testgen
-    testgen_code_markdown = extract_code_string_context_from_files(
+    testgen_context = extract_code_markdown_context_from_files(
         helpers_of_fto_dict,
         helpers_of_helpers_dict,
         project_root_path,
         remove_docstrings=False,
         code_context_type=CodeContextType.TESTGEN,
     )
-    testgen_context_code = testgen_code_markdown.code
-    testgen_context_code_tokens = encoded_tokens_len(testgen_context_code)
-    if testgen_context_code_tokens > testgen_token_limit:
-        testgen_code_markdown = extract_code_string_context_from_files(
+    testgen_markdown_code = testgen_context.markdown
+    testgen_code_token_length = encoded_tokens_len(testgen_markdown_code)
+    if testgen_code_token_length > testgen_token_limit:
+        testgen_context = extract_code_markdown_context_from_files(
             helpers_of_fto_dict,
             helpers_of_helpers_dict,
             project_root_path,
             remove_docstrings=True,
             code_context_type=CodeContextType.TESTGEN,
         )
-        testgen_context_code = testgen_code_markdown.code
-        testgen_context_code_tokens = encoded_tokens_len(testgen_context_code)
-        if testgen_context_code_tokens > testgen_token_limit:
+        testgen_markdown_code = testgen_context.markdown
+        testgen_code_token_length = encoded_tokens_len(testgen_markdown_code)
+        if testgen_code_token_length > testgen_token_limit:
             raise ValueError("Testgen code context has exceeded token limit, cannot proceed")
     code_hash_context = hashing_code_context.markdown
     code_hash = hashlib.sha256(code_hash_context.encode("utf-8")).hexdigest()
 
     return CodeOptimizationContext(
-        testgen_context_code=testgen_context_code,
+        testgen_context=testgen_context,
         read_writable_code=final_read_writable_code,
         read_only_context_code=read_only_context_code,
         hashing_code_context=code_hash_context,

codeflash/discovery/pytest_new_process_discovery.py

@@ -1,6 +1,9 @@
 # ruff: noqa
 import sys
+from pathlib import Path
 from typing import Any
+import pickle
+
 
 # This script should not have any relation to the codeflash package, be careful with imports
 cwd = sys.argv[1]
@@ -11,44 +14,48 @@
 sys.path.insert(1, str(cwd))
 
 
+def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]:
+    test_results = []
+    for test in pytest_tests:
+        test_class = None
+        if test.cls:
+            test_class = test.parent.name
+        test_results.append({"test_file": str(test.path), "test_class": test_class, "test_function": test.name})
+    return test_results
+
+
 class PytestCollectionPlugin:
     def pytest_collection_finish(self, session) -> None:
-        global pytest_rootdir
+        global pytest_rootdir, collected_tests
+
         collected_tests.extend(session.items)
         pytest_rootdir = session.config.rootdir
 
+        # Write results immediately since pytest.main() will exit after this callback, not always with a success code
+        tests = parse_pytest_collection_results(collected_tests)
+        exit_code = getattr(session.config, "exitstatus", 0)
+        with Path(pickle_path).open("wb") as f:
+            pickle.dump((exit_code, tests, pytest_rootdir), f, protocol=pickle.HIGHEST_PROTOCOL)
+
     def pytest_collection_modifyitems(self, items) -> None:
         skip_benchmark = pytest.mark.skip(reason="Skipping benchmark tests")
         for item in items:
             if "benchmark" in item.fixturenames:
                 item.add_marker(skip_benchmark)
 
 
-def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]:
-    test_results = []
-    for test in pytest_tests:
-        test_class = None
-        if test.cls:
-            test_class = test.parent.name
-        test_results.append({"test_file": str(test.path), "test_class": test_class, "test_function": test.name})
-    return test_results
-
-
 if __name__ == "__main__":
-    from pathlib import Path
-
     import pytest
 
     try:
-        exitcode = pytest.main(
-            [tests_root, "-p no:logging", "--collect-only", "-m", "not skip", "-p", "no:codeflash-benchmark"],
+        pytest.main(
+            [tests_root, "-p", "no:logging", "--collect-only", "-m", "not skip", "-p", "no:codeflash-benchmark"],
             plugins=[PytestCollectionPlugin()],
         )
     except Exception as e:
         print(f"Failed to collect tests: {e!s}")
-        exitcode = -1
-    tests = parse_pytest_collection_results(collected_tests)
-    import pickle
-
-    with Path(pickle_path).open("wb") as f:
-        pickle.dump((exitcode, tests, pytest_rootdir), f, protocol=pickle.HIGHEST_PROTOCOL)
+        try:
+            with Path(pickle_path).open("wb") as f:
+                pickle.dump((-1, [], None), f, protocol=pickle.HIGHEST_PROTOCOL)
+        except Exception as pickle_error:
+            print(f"Failed to write failure pickle: {pickle_error!s}", file=sys.stderr)

codeflash/models/models.py

@@ -163,7 +163,7 @@ class CodeString(BaseModel):
 
 
 def get_code_block_splitter(file_path: Path) -> str:
-    return f"# file: {file_path}"
+    return f"# file: {file_path.as_posix()}"
 
 
 markdown_pattern = re.compile(r"```python:([^\n]+)\n(.*?)\n```", re.DOTALL)
@@ -254,7 +254,7 @@ def parse_markdown_code(markdown_code: str) -> CodeStringsMarkdown:
 
 
 class CodeOptimizationContext(BaseModel):
-    testgen_context_code: str = ""
+    testgen_context: CodeStringsMarkdown
     read_writable_code: CodeStringsMarkdown
     read_only_context_code: str = ""
     hashing_code_context: str = ""

codeflash/optimization/function_optimizer.py

@@ -309,7 +309,7 @@ def generate_and_instrument_tests(
             revert_to_print=bool(get_pr_number()),
         ):
             generated_results = self.generate_tests_and_optimizations(
-                testgen_context_code=code_context.testgen_context_code,
+                testgen_context=code_context.testgen_context,
                 read_writable_code=code_context.read_writable_code,
                 read_only_context_code=code_context.read_only_context_code,
                 helper_functions=code_context.helper_functions,
@@ -345,7 +345,6 @@ def generate_and_instrument_tests(
             logger.info(f"Generated test {i + 1}/{count_tests}:")
             code_print(generated_test.generated_original_test_source, file_name=f"test_{i + 1}.py")
         if concolic_test_str:
-            # no concolic tests in lsp mode
             logger.info(f"Generated test {count_tests}/{count_tests}:")
             code_print(concolic_test_str)
 
@@ -972,7 +971,7 @@ def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]:
 
         return Success(
             CodeOptimizationContext(
-                testgen_context_code=new_code_ctx.testgen_context_code,
+                testgen_context=new_code_ctx.testgen_context,
                 read_writable_code=new_code_ctx.read_writable_code,
                 read_only_context_code=new_code_ctx.read_only_context_code,
                 hashing_code_context=new_code_ctx.hashing_code_context,
@@ -1079,7 +1078,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
 
     def generate_tests_and_optimizations(
         self,
-        testgen_context_code: str,
+        testgen_context: CodeStringsMarkdown,
         read_writable_code: CodeStringsMarkdown,
         read_only_context_code: str,
         helper_functions: list[FunctionSource],
@@ -1093,7 +1092,7 @@ def generate_tests_and_optimizations(
         # Submit the test generation task as future
         future_tests = self.submit_test_generation_tasks(
             self.executor,
-            testgen_context_code,
+            testgen_context.markdown,
             [definition.fully_qualified_name for definition in helper_functions],
             generated_test_paths,
             generated_perf_test_paths,

tests/test_code_replacement.py

@@ -798,7 +798,8 @@ def main_method(self):
 
 
 def test_code_replacement10() -> None:
-    get_code_output = """from __future__ import annotations
+    get_code_output = """# file: test_code_replacement.py
+from __future__ import annotations
 
 class HelperClass:
     def __init__(self, name):
@@ -828,7 +829,7 @@ def main_method(self):
     )
     func_optimizer = FunctionOptimizer(function_to_optimize=func_top_optimize, test_cfg=test_config)
     code_context = func_optimizer.get_code_optimization_context().unwrap()
-    assert code_context.testgen_context_code.rstrip() == get_code_output.rstrip()
+    assert code_context.testgen_context.flat.rstrip() == get_code_output.rstrip()
 
 
 def test_code_replacement11() -> None:

tests/test_code_utils.py

@@ -21,6 +21,7 @@
 )
 from codeflash.code_utils.concolic_utils import clean_concolic_tests
 from codeflash.code_utils.coverage_utils import extract_dependent_function, generate_candidates, prepare_coverage_files
+from codeflash.models.models import CodeStringsMarkdown
 
 
 @pytest.fixture
@@ -382,69 +383,76 @@ def mock_code_context():
 def test_extract_dependent_function_sync_and_async(mock_code_context):
     """Test extract_dependent_function with both sync and async functions."""
     # Test sync function extraction
-    mock_code_context.testgen_context_code = """
+    mock_code_context.testgen_context = CodeStringsMarkdown.parse_markdown_code("""```python:file.py
 def main_function():
     pass
 
 def helper_function():
     pass
-"""
+```
+""")
     assert extract_dependent_function("main_function", mock_code_context) == "helper_function"
 
     # Test async function extraction
-    mock_code_context.testgen_context_code = """
+    mock_code_context.testgen_context = CodeStringsMarkdown.parse_markdown_code("""```python:file.py
 def main_function():
     pass
 
 async def async_helper_function():
     pass
-"""
+```
+""")
+
     assert extract_dependent_function("main_function", mock_code_context) == "async_helper_function"
 
 
 def test_extract_dependent_function_edge_cases(mock_code_context):
     """Test extract_dependent_function edge cases."""
     # No dependent functions
-    mock_code_context.testgen_context_code = """
+    mock_code_context.testgen_context = CodeStringsMarkdown.parse_markdown_code("""```python:file.py
 def main_function():
     pass
-"""
+```
+""")
     assert extract_dependent_function("main_function", mock_code_context) is False
 
     # Multiple dependent functions
-    mock_code_context.testgen_context_code = """
+    mock_code_context.testgen_context = CodeStringsMarkdown.parse_markdown_code("""```python:file.py
 def main_function():
     pass
-
 def helper1():
     pass
 
 async def helper2():
     pass
-"""
+```
+""")
     assert extract_dependent_function("main_function", mock_code_context) is False
 
 
 def test_extract_dependent_function_mixed_scenarios(mock_code_context):
     """Test extract_dependent_function with mixed sync/async scenarios."""
     # Async main with sync helper
-    mock_code_context.testgen_context_code = """
+    mock_code_context.testgen_context = CodeStringsMarkdown.parse_markdown_code("""```python:file.py
 async def async_main():
     pass
 
 def sync_helper():
     pass
-"""
+```
+""")
     assert extract_dependent_function("async_main", mock_code_context) == "sync_helper"
 
     # Only async functions
-    mock_code_context.testgen_context_code = """
+    mock_code_context.testgen_context = CodeStringsMarkdown.parse_markdown_code("""```python:file.py
 async def async_main():
     pass
 
 async def async_helper():
     pass
-"""
+```
+""")
+
     assert extract_dependent_function("async_main", mock_code_context) == "async_helper"
 
 

tests/test_function_dependencies.py

@@ -160,8 +160,9 @@ def test_class_method_dependencies() -> None:
     )
     assert code_context.helper_functions[0].qualified_name == "Graph.topologicalSortUtil"
     assert (
-        code_context.testgen_context_code
-        == """from collections import defaultdict
+        code_context.testgen_context.flat
+        == """# file: test_function_dependencies.py
+from collections import defaultdict
 
 class Graph:
     def __init__(self, vertices):
@@ -220,8 +221,9 @@ def test_recursive_function_context() -> None:
     assert code_context.helper_functions[0].fully_qualified_name == "test_function_dependencies.C.calculate_something_3"
     assert code_context.helper_functions[1].fully_qualified_name == "test_function_dependencies.C.recursive"
     assert (
-        code_context.testgen_context_code
-        == """class C:
+        code_context.testgen_context.flat
+        == """# file: test_function_dependencies.py
+class C:
     def calculate_something_3(self, num):
         return num + 1