changes after code review, made logic more robust, added tests_root arg to codeflash capture, added recursive tests, renamed some files / functions

Author: alvin-r

codeflash/optimization/optimizer.py

@@ -69,7 +69,7 @@
 from codeflash.verification.bayesian_analysis import compare_function_runtime_distributions
 from codeflash.verification.concolic_testing import generate_concolic_tests
 from codeflash.verification.equivalence import compare_test_results
-from codeflash.verification.instrument_code import instrument_code
+from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
 from codeflash.verification.parse_test_output import parse_test_results
 from codeflash.verification.test_results import TestResults, TestType
 from codeflash.verification.test_runner import run_behavioral_tests, run_benchmarking_tests
@@ -1029,7 +1029,9 @@ def establish_original_code_baseline(
             original_fto_code = function_file_path.read_text("utf-8")
             # Instrument codeflash capture
             try:
-                instrument_code(function_to_optimize, file_path_to_helper_classes)
+                instrument_codeflash_capture(
+                    function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root
+                )
                 behavioral_results, coverage_results = self.run_and_parse_tests(
                     testing_type=TestingMode.BEHAVIOR,
                     test_env=test_env,
@@ -1175,7 +1177,9 @@ def run_optimized_candidate(
             for module_abspath in original_helper_code:
                 candidate_helper_code[module_abspath] = Path(module_abspath).read_text("utf-8")
             try:
-                instrument_code(function_to_optimize, file_path_to_helper_classes)
+                instrument_codeflash_capture(
+                    function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root
+                )
 
                 candidate_behavior_results, _ = self.run_and_parse_tests(
                     testing_type=TestingMode.BEHAVIOR,

codeflash/verification/codeflash_capture.py

@@ -6,55 +6,81 @@
 import os
 import sqlite3
 import time
+from pathlib import Path
 
 import dill as pickle
 
 from codeflash.verification.test_results import VerificationType
 
 
-def get_test_info_from_stack() -> tuple[str, str | None, str, str]:
-    """Extract test information from the call stack."""
-    stack = inspect.stack()
-
-    # Default values
+def get_test_info_from_stack(tests_root: str) -> tuple[str, str | None, str, str]:
+    """Extract test information by walking the call stack from the current frame."""
     test_module_name = ""
-    test_class_name = None
-    test_name = None
-    line_id = ""  # Note that the way this line_id is defined is from the line_id called in instrumentation
-
-    # Search through stack for test information
-    for frame in stack:
-        if frame.function.startswith("test_"):  # May need a more robust way to find the test file
-            test_name = frame.function
-            test_module_name = inspect.getmodule(frame[0]).__name__
-            line_id = str(frame.lineno)
+    test_class_name: str | None = None
+    test_name: str | None = None
+    line_id = ""
+
+    # Get current frame and skip our own function's frame
+    frame = inspect.currentframe()
+    if frame is not None:
+        frame = frame.f_back
+
+    # Walk the stack
+    while frame is not None:
+        function_name = frame.f_code.co_name
+        filename = frame.f_code.co_filename
+        lineno = frame.f_lineno
+
+        # Check if function name indicates a test (e.g., starts with "test_")
+        if function_name.startswith("test_"):
+            test_name = function_name
+            test_module = inspect.getmodule(frame)
+            if hasattr(test_module, "__name__"):
+                test_module_name = test_module.__name__
+            line_id = str(lineno)
+
             # Check if it's a method in a class
-            if "self" in frame.frame.f_locals:
-                test_class_name = frame.frame.f_locals["self"].__class__.__name__
+            if (
+                "self" in frame.f_locals
+                and hasattr(frame.f_locals["self"], "__class__")
+                and hasattr(frame.f_locals["self"].__class__, "__name__")
+            ):
+                test_class_name = frame.f_locals["self"].__class__.__name__
             break
-        # Check if module name starts with test
-        module_name = frame.frame.f_globals["__name__"]
-        if module_name and module_name.split(".")[-1].startswith("test_"):
-            test_module_name = module_name
-            line_id = str(frame.lineno)
-            if frame.function != "<module>":
-                test_name = frame.function  # Technically not a test, but save the info since there is no test function
-            # Check if it's in a class
-            if "self" in frame.frame.f_locals:
-                test_class_name = frame.frame.f_locals["self"].__class__.__name__
+
+        # Check for instantiation on the module level
+        if (
+            "__name__" in frame.f_globals
+            and frame.f_globals["__name__"].split(".")[-1].startswith("test_")
+            and Path(filename).resolve().is_relative_to(Path(tests_root))
+            and function_name == "<module>"
+        ):
+            test_module_name = frame.f_globals["__name__"]
+            line_id = str(lineno)
+
+            #     # Check if it's a method in a class
+            if (
+                "self" in frame.f_locals
+                and hasattr(frame.f_locals["self"], "__class__")
+                and hasattr(frame.f_locals["self"].__class__, "__name__")
+            ):
+                test_class_name = frame.f_locals["self"].__class__.__name__
             break
 
+        # Go to the previous frame
+        frame = frame.f_back
+
     return test_module_name, test_class_name, test_name, line_id
 
 
-def codeflash_capture(function_name: str, tmp_dir_path: str, is_fto: bool = False):
+def codeflash_capture(function_name: str, tmp_dir_path: str, tests_root: str, is_fto: bool = False):
     """Defines decorator to be instrumented onto the init function in the code. Collects info of the test that called this, and captures the state of the instance."""
 
     def decorator(wrapped):
         @functools.wraps(wrapped)
         def wrapper(*args, **kwargs):
             # Dynamic information retrieved from stack
-            test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack()
+            test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(tests_root)
 
             # Get env variables
             loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
@@ -98,7 +124,12 @@ def wrapper(*args, **kwargs):
                 gc.enable()
 
             # Capture instance state after initialization
-            instance_state = args[0].__dict__  # self is always the first argument
+            if hasattr(args[0], "__dict__"):
+                instance_state = args[
+                    0
+                ].__dict__  # self is always the first argument, this is ensured during instrumentation
+            else:
+                raise ValueError("Instance state could not be captured.")
             codeflash_cur.execute(
                 "CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT)"
             )

codeflash/verification/comparator.py

@@ -51,6 +51,7 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool:
         if type(orig) is not type(new):
             type_obj = type(orig)
             new_type_obj = type(new)
+            # distinct type objects are created at runtime, even if the class code is exactly the same, so we can only compare the names
             if type_obj.__name__ != new_type_obj.__name__ or type_obj.__qualname__ != new_type_obj.__qualname__:
                 return False
         if isinstance(orig, (list, tuple)):

codeflash/verification/equivalence.py

@@ -46,14 +46,19 @@ def compare_test_results(original_results: TestResults, candidate_results: TestR
         if not comparator(original_test_result.return_value, cdd_test_result.return_value, superset_obj=superset_obj):
             are_equal = False
             logger.debug(
-                f"""
-                File Name: {original_test_result.file_name}
-                Test Type: {original_test_result.test_type}
-                Verification Type: {original_test_result.verification_type}
-                Invocation ID: {original_test_result.id}
-                Original return value: {original_test_result.return_value}
-                CDD return value: {cdd_test_result.return_value}
-                -------------------"""
+                "File Name: %s\n"
+                "Test Type: %s\n"
+                "Verification Type: %s\n"
+                "Invocation ID: %s\n"
+                "Original return value: %s\n"
+                "Candidate return value: %s\n"
+                "-------------------",
+                original_test_result.file_name,
+                original_test_result.test_type,
+                original_test_result.verification_type,
+                original_test_result.id,
+                original_test_result.return_value,
+                cdd_test_result.return_value,
             )
             break
         if original_test_result.test_type in [TestType.EXISTING_UNIT_TEST, TestType.CONCOLIC_COVERAGE_TEST] and (

codeflash/verification/instrument_code.py …fication/instrument_codeflash_capture.pycodeflash/verification/instrument_code.py renamed to codeflash/verification/instrument_codeflash_capture.py codeflash/verification/instrument_code.py renamed to codeflash/verification/instrument_codeflash_capture.py

@@ -9,7 +9,9 @@
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 
 
-def instrument_code(function_to_optimize: FunctionToOptimize, file_path_to_helper_class: dict[Path, set[str]]) -> None:
+def instrument_codeflash_capture(
+    function_to_optimize: FunctionToOptimize, file_path_to_helper_class: dict[Path, set[str]], tests_root: Path
+) -> None:
     """Instrument __init__ function with codeflash_capture decorator if it's in a class."""
     # Find the class parent
     if len(function_to_optimize.parents) == 1 and function_to_optimize.parents[0].type == "ClassDef":
@@ -23,42 +25,38 @@ def instrument_code(function_to_optimize: FunctionToOptimize, file_path_to_helpe
     ):
         file_path_to_helper_class[function_to_optimize.file_path].remove(class_parent.name)
     # Instrument fto class
-    with open(function_to_optimize.file_path) as f:
-        original_code = f.read()
-
+    original_code = function_to_optimize.file_path.read_text(encoding="utf-8")
     # Add decorator to init
     modified_code = add_codeflash_capture_to_init(
         target_classes={class_parent.name},
         fto_name=function_to_optimize.function_name,
         tmp_dir_path=str(get_run_tmp_file(Path("test_return_values"))),
         code=original_code,
+        tests_root=tests_root,
         is_fto=True,
     )
-
-    with open(function_to_optimize.file_path, "w") as f:
-        f.write(modified_code)
+    function_to_optimize.file_path.write_text(modified_code, encoding="utf-8")
 
     # Instrument helper classes
     for file_path, helper_classes in file_path_to_helper_class.items():
-        with open(file_path) as f:
-            original_code = f.read()
+        original_code = file_path.read_text(encoding="utf-8")
         modified_code = add_codeflash_capture_to_init(
             target_classes=helper_classes,
             fto_name=function_to_optimize.function_name,
             tmp_dir_path=str(get_run_tmp_file(Path("test_return_values"))),
             code=original_code,
+            tests_root=tests_root,
             is_fto=False,
         )
-        with open(file_path, "w") as f:
-            f.write(modified_code)
+        file_path.write_text(modified_code, encoding="utf-8")
 
 
 def add_codeflash_capture_to_init(
-    target_classes: set[str], fto_name: str, tmp_dir_path: str, code: str, is_fto: bool = False
+    target_classes: set[str], fto_name: str, tmp_dir_path: str, code: str, tests_root: Path, is_fto: bool = False
 ) -> str:
     """Add codeflash_capture decorator to __init__ function in the specified class."""
     tree = ast.parse(code)
-    transformer = InitDecorator(target_classes, fto_name, tmp_dir_path, is_fto)
+    transformer = InitDecorator(target_classes, fto_name, tmp_dir_path, tests_root, is_fto)
     modified_tree = transformer.visit(tree)
     if transformer.inserted_decorator:
         ast.fix_missing_locations(modified_tree)
@@ -70,12 +68,15 @@ def add_codeflash_capture_to_init(
 class InitDecorator(ast.NodeTransformer):
     """AST transformer that adds codeflash_capture decorator to specific class's __init__."""
 
-    def __init__(self, target_classes: set[str], fto_name: str, tmp_dir_path: str, is_fto=False) -> None:
+    def __init__(
+        self, target_classes: set[str], fto_name: str, tmp_dir_path: str, tests_root: Path, is_fto=False
+    ) -> None:
         self.target_classes = target_classes
         self.fto_name = fto_name
         self.tmp_dir_path = tmp_dir_path
         self.is_fto = is_fto
         self.has_import = False
+        self.tests_root = tests_root
         self.inserted_decorator = False
 
     def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.ImportFrom:
@@ -110,12 +111,19 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
             keywords=[
                 ast.keyword(arg="function_name", value=ast.Constant(value=".".join([node.name, "__init__"]))),
                 ast.keyword(arg="tmp_dir_path", value=ast.Constant(value=self.tmp_dir_path)),
+                ast.keyword(arg="tests_root", value=ast.Constant(value=str(self.tests_root))),
                 ast.keyword(arg="is_fto", value=ast.Constant(value=self.is_fto)),
             ],
         )
 
         for item in node.body:
-            if isinstance(item, ast.FunctionDef) and item.name == "__init__":
+            if (
+                isinstance(item, ast.FunctionDef)
+                and item.name == "__init__"
+                and item.args.args
+                and isinstance(item.args.args[0], ast.arg)
+                and item.args.args[0].arg == "self"
+            ):
                 has_init = True
 
                 # Add decorator at the start of the list if not already present

codeflash/verification/parse_test_output.py

@@ -412,7 +412,9 @@ def merge_test_results(
                         test_type=xml_result.test_type,
                         return_value=result_bin.return_value,
                         timed_out=xml_result.timed_out,
-                        verification_type=VerificationType(result_bin.verification_type),
+                        verification_type=VerificationType(result_bin.verification_type)
+                        if result_bin.verification_type
+                        else None,
                     )
                 )
         elif xml_results.test_results[0].id.iteration_id is not None:
@@ -439,7 +441,9 @@ def merge_test_results(
                         timed_out=xml_result.timed_out
                         if bin_result.runtime is None
                         else False,  # If runtime was measured in the bin file, then the testcase did not time out
-                        verification_type=VerificationType(bin_result.verification_type),
+                        verification_type=VerificationType(bin_result.verification_type)
+                        if bin_result.verification_type
+                        else None,
                     )
                 )
         else:
@@ -463,7 +467,9 @@ def merge_test_results(
                         test_type=bin_result.test_type,
                         return_value=bin_result.return_value,
                         timed_out=xml_result.timed_out,  # only the xml gets the timed_out flag
-                        verification_type=VerificationType(bin_result.verification_type),
+                        verification_type=VerificationType(bin_result.verification_type)
+                        if bin_result.verification_type
+                        else None,
                     )
                 )