cleanup

Author: KRRT7

codeflash/code_utils/code_replacer.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import ast
+import re
 from collections import defaultdict
 from functools import lru_cache
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Optional, TypeVar
 
 import libcst as cst
 
@@ -338,25 +339,91 @@ def function_to_optimize_original_worktree_fqn(
     )
 
 
+class AssertCleanup:
+    def transform_asserts(self, code: str) -> str:
+        lines = code.splitlines()
+        result_lines = []
+
+        for line in lines:
+            transformed = self._transform_assert_line(line)
+            if transformed is not None:
+                result_lines.append(transformed)
+            else:
+                result_lines.append(line)
+
+        return "\n".join(result_lines)
+
+    def _transform_assert_line(self, line: str) -> Optional[str]:
+        indent = line[: len(line) - len(line.lstrip())]
+
+        assert_match = re.match(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$", line)
+        if assert_match:
+            expression = assert_match.group(1).strip()
+            if expression.startswith("not "):
+                return f"{indent}{expression}"
+
+            expression = re.sub(r"[,;]\s*$", "", expression)
+            return f"{indent}{expression}"
+
+        unittest_match = re.match(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$", line)
+        if unittest_match:
+            indent, assert_method, args = unittest_match.groups()
+
+            if args:
+                arg_parts = self._split_top_level_args(args)
+                if arg_parts and arg_parts[0]:
+                    return f"{indent}{arg_parts[0]}"
+
+        return None
+
+    def _split_top_level_args(self, args_str: str) -> list[str]:
+        result = []
+        current = []
+        depth = 0
+
+        for char in args_str:
+            if char in "([{":
+                depth += 1
+                current.append(char)
+            elif char in ")]}":
+                depth -= 1
+                current.append(char)
+            elif char == "," and depth == 0:
+                result.append("".join(current).strip())
+                current = []
+            else:
+                current.append(char)
+
+        if current:
+            result.append("".join(current).strip())
+
+        return result
+
+
 def clean_concolic_tests(test_suite_code: str) -> str:
     try:
+        can_parse = True
         tree = ast.parse(test_suite_code)
+    except SyntaxError:
+        can_parse = False
+
+    if not can_parse:
+        return AssertCleanup().transform_asserts(test_suite_code)
 
-        for node in ast.walk(tree):
-            if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
-                new_body = []
-                for stmt in node.body:
-                    if isinstance(stmt, ast.Assert):
-                        if isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Call):
-                            new_body.append(ast.Expr(value=stmt.test.left))
-                        else:
-                            new_body.append(stmt)
+    tree = ast.parse(test_suite_code)
 
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
+            new_body = []
+            for stmt in node.body:
+                if isinstance(stmt, ast.Assert):
+                    if isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Call):
+                        new_body.append(ast.Expr(value=stmt.test.left))
                     else:
                         new_body.append(stmt)
-                node.body = new_body
 
-        return ast.unparse(tree).strip()
-    except SyntaxError:
-        logger.warning("Failed to parse and modify CrossHair generated tests.  Using original output.")
-        return test_suite_code
+                else:
+                    new_body.append(stmt)
+            node.body = new_body
+
+    return ast.unparse(tree).strip()

codeflash/verification/concolic_testing.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import ast
-import difflib
 import subprocess
 import tempfile
 from argparse import Namespace
@@ -60,8 +59,8 @@ def generate_concolic_tests(
             return function_to_concolic_tests, concolic_test_suite_code
 
         if cover_result.returncode == 0:
-            original_code: str = cover_result.stdout
-            concolic_test_suite_code: str = clean_concolic_tests(original_code)
+            generated_concolic_test: str = cover_result.stdout
+            concolic_test_suite_code: str = clean_concolic_tests(generated_concolic_test)
             concolic_test_suite_dir = Path(tempfile.mkdtemp(dir=test_cfg.concolic_test_root_dir))
             concolic_test_suite_path = concolic_test_suite_dir / "test_concolic_coverage.py"
             concolic_test_suite_path.write_text(concolic_test_suite_code, encoding="utf8")

tests/test_code_utils.py

@@ -427,3 +427,17 @@ def test_calculate_tuple_sum(my_tuple: tuple[int, int, int]) -> None:
     calculate_tuple_sum((1, 2, 3))
 """
     assert cleaned_code == expected_cleaned_code.strip()
+
+    concolic_generated_repr_code = """from src.blib2to3.pgen2.grammar import Grammar
+
+def test_Grammar_copy():
+    assert Grammar.copy(Grammar()) == <src.blib2to3.pgen2.grammar.Grammar object at 0x104c30f50>
+"""
+    cleaned_code = clean_concolic_tests(concolic_generated_repr_code)
+    expected_cleaned_code = """
+from src.blib2to3.pgen2.grammar import Grammar
+
+def test_Grammar_copy():
+    Grammar.copy(Grammar())
+"""
+    assert cleaned_code == expected_cleaned_code.strip()