⚡️ Speed up method AssertCleanup.transform_asserts by 46% in PR #26 (clean_concolic_tests)

codeflash/code_utils/concolic_utils.py

@@ -12,26 +12,23 @@ def transform_asserts(self, code: str) -> str:
 
         for line in lines:
             transformed = self._transform_assert_line(line)
-            if transformed is not None:
-                result_lines.append(transformed)
-            else:
-                result_lines.append(line)
+            result_lines.append(transformed if transformed is not None else 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)
+        assert_match = self.assert_re.match(line)
         if assert_match:
             expression = assert_match.group(1).strip()
             if expression.startswith("not "):
                 return f"{indent}{expression}"
 
-            expression = re.sub(r"[,;]\s*$", "", expression)
+            expression = expression.rstrip(",;")
             return f"{indent}{expression}"
 
-        unittest_match = re.match(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$", line)
+        unittest_match = self.unittest_re.match(line)
         if unittest_match:
             indent, assert_method, args = unittest_match.groups()
 
@@ -65,6 +62,11 @@ def _split_top_level_args(self, args_str: str) -> list[str]:
 
         return result
 
+    def __init__(self):
+        # Pre-compiling regular expressions for faster execution
+        self.assert_re = re.compile(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$")
+        self.unittest_re = re.compile(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$")
+
 
 def clean_concolic_tests(test_suite_code: str) -> str:
     try: