some progress

Author: aseembits93

code_to_optimize/tests/pytest/conftest.py

@@ -1,26 +1,17 @@
 import pytest
 import time
 
-# @pytest.fixture(autouse=True)
-# def fixture(request):
-#     if request.node.get_closest_marker("no_autouse"):
-#         # Skip the fixture logic
-#         yield
-#     else:
-#         start_time = time.time()
-#         time.sleep(0.1)
-#         yield
-#         print(f"Took {time.time() - start_time} seconds")
-
 
 @pytest.fixture(autouse=True)
-def fixture1(request):  # We don't need this fixture during testing
+def fixture1(request):
     start_time = time.time()
     time.sleep(0.1)
     yield
     print(f"Took {time.time() - start_time} seconds")
 
 
 @pytest.fixture(autouse=True)
-def fixture2(request):  # We need it
+def fixture2(request):  # We don't need this fixture during testing
+    print("not doing anything")
     yield
+    print("did nothing")

code_to_optimize/tests/pytest/test_bubble_sort.py

@@ -1,7 +1,6 @@
 from code_to_optimize.bubble_sort import sorter
-import pytest
 
-@pytest.mark.no_autouse
+
 def test_sort():
     input = [5, 4, 3, 2, 1, 0]
     output = sorter(input)

codeflash/cli_cmds/cmd_init.py

@@ -34,7 +34,7 @@
     from argparse import Namespace
 
 CODEFLASH_LOGO: str = (
-    f"{LF}"
+    f"{LF}"  # noqa : ISC003
     r"                   _          ___  _               _     " + f"{LF}"
     r"                  | |        / __)| |             | |    " + f"{LF}"
     r"  ____   ___    _ | |  ____ | |__ | |  ____   ___ | | _  " + f"{LF}"

codeflash/code_utils/code_extractor.py

@@ -5,9 +5,11 @@
 from typing import TYPE_CHECKING, Optional
 
 import libcst as cst
+from libcst import MetadataWrapper
 from libcst.codemod import CodemodContext
 from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, RemoveImportsVisitor
 from libcst.helpers import calculate_module_and_package
+from libcst.metadata import FullyQualifiedNameProvider
 
 from codeflash.cli_cmds.console import logger
 from codeflash.models.models import FunctionParent
@@ -21,6 +23,53 @@
     from codeflash.models.models import FunctionSource
 
 
+class FunctionNameCollector(cst.CSTVisitor):
+    """A LibCST visitor that collects the fully qualified names of all functions."""
+
+    METADATA_DEPENDENCIES = (FullyQualifiedNameProvider,)
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.qualified_names: set[str] = set()
+
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
+        """Visits a function definition node and extracts its qualified name."""
+        try:
+            q_names = self.get_metadata(FullyQualifiedNameProvider, node)
+            for q_name in q_names:
+                self.qualified_names.add(q_name.name)
+        except KeyError:
+            # This can happen for functions defined in scopes where a qualified
+            # name cannot be determined.
+            pass
+
+
+def get_function_qualified_names(file_path: Path) -> list[str]:
+    """Parse a Python file and returns a list of fully qualified function names.
+
+    Args:
+        file_path: The path to the Python file.
+
+    Returns:
+        A list of string representations of the qualified function names.
+
+    """
+    with file_path.open("r") as f:
+        source_code = f.read()
+
+    # Parse the source code into a CST
+    module = cst.parse_module(source_code)
+
+    # Wrap the module with a metadata wrapper to enable name resolution
+    wrapper = MetadataWrapper(module)
+
+    # Create an instance of the visitor and visit the wrapped module
+    visitor = FunctionNameCollector()
+    wrapper.visit(visitor)
+
+    return list(visitor.qualified_names)
+
+
 class GlobalAssignmentCollector(cst.CSTVisitor):
     """Collects all global assignment statements."""
 

codeflash/code_utils/code_replacer.py

@@ -1,14 +1,23 @@
 from __future__ import annotations
 
 import ast
+import contextlib
 from collections import defaultdict
 from functools import lru_cache
 from typing import TYPE_CHECKING, Optional, TypeVar
 
+import isort
 import libcst as cst
+import libcst.matchers as m
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.code_extractor import add_global_assignments, add_needed_imports_from_module
+from codeflash.code_utils.code_extractor import (
+    add_global_assignments,
+    add_needed_imports_from_module,
+    get_function_qualified_names,
+)
+from codeflash.code_utils.config_parser import find_conftest_files
+from codeflash.code_utils.line_profile_utils import ImportAdder, add_decorator_to_qualified_function
 from codeflash.models.models import FunctionParent
 
 if TYPE_CHECKING:
@@ -33,18 +42,78 @@ def normalize_code(code: str) -> str:
     return ast.unparse(normalize_node(ast.parse(code)))
 
 
-# def modify_autouse_fixture():
-#     # find fixutre definition in conftetst.py (the one closest to the test)
-#     # get fixtures present in override-fixtures in pyproject.toml
-#     # add if marker closest return
-#     conftest_files = find_conftest_files()
-#     for cf_file in conftest_files:
-#         # iterate over all functions in the file
-#         # if function has autouse fixture, modify function to bypass with custom marker
-#         pass
+class AutouseFixtureModifier(cst.CSTTransformer):
+    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
+        # Matcher for '@fixture' or '@pytest.fixture'
+        fixture_decorator_func = m.Name("fixture") | m.Attribute(value=m.Name("pytest"), attr=m.Name("fixture"))
+
+        for decorator in original_node.decorators:
+            if m.matches(
+                decorator,
+                m.Decorator(
+                    decorator=m.Call(
+                        func=fixture_decorator_func, args=[m.Arg(value=m.Name("True"), keyword=m.Name("autouse"))]
+                    )
+                ),
+            ):
+                # Found a matching fixture with autouse=True
+
+                # 1. The original body of the function will become the 'else' block.
+                #    updated_node.body is an IndentedBlock, which is what cst.Else expects.
+                else_block = cst.Else(body=updated_node.body)
 
+                # 2. Create the new 'if' block that will exit the fixture early.
+                if_test = cst.parse_expression('request.node.get_closest_marker("no_autouse")')
+                yield_statement = cst.parse_statement("yield")
+                if_body = cst.IndentedBlock(body=[yield_statement])
 
-# reuse line profiler utils to add decorator and import to test fns
+                # 3. Construct the full if/else statement.
+                new_if_statement = cst.If(test=if_test, body=if_body, orelse=else_block)
+
+                # 4. Replace the entire function's body with our new single statement.
+                return updated_node.with_changes(body=cst.IndentedBlock(body=[new_if_statement]))
+        return updated_node
+
+
+@contextlib.contextmanager
+def disable_autouse(test_path: Path) -> None:
+    file_content = test_path.read_text(encoding="utf-8")
+    try:
+        module = cst.parse_module(file_content)
+        disable_autouse_fixture = AutouseFixtureModifier()
+        modified_module = module.visit(disable_autouse_fixture)
+        test_path.write_text(modified_module.code, encoding="utf-8")
+    finally:
+        test_path.write_text(file_content, encoding="utf-8")
+
+
+def modify_autouse_fixture(test_paths: list[Path]) -> None:
+    # find fixutre definition in conftetst.py (the one closest to the test)
+    # get fixtures present in override-fixtures in pyproject.toml
+    # add if marker closest return
+    conftest_files = find_conftest_files(test_paths)
+    for cf_file in conftest_files:
+        # iterate over all functions in the file
+        # if function has autouse fixture, modify function to bypass with custom marker
+        disable_autouse(cf_file)
+
+
+# # reuse line profiler utils to add decorator and import to test fns
+def add_custom_marker_to_all_tests(test_paths: list[Path]) -> None:
+    for test_path in test_paths:
+        # read file
+        file_content = test_path.read_text(encoding="utf-8")
+        module = cst.parse_module(file_content)
+        importadder = ImportAdder("import pytest")
+        modified_module = module.visit(importadder)
+        modified_module = isort.code(modified_module.code, float_to_top=True)
+        qualified_fn_names = get_function_qualified_names(test_path)
+        for fn_name in qualified_fn_names:
+            modified_module = add_decorator_to_qualified_function(
+                modified_module, fn_name, "pytest.mark.codeflash_no_autouse"
+            )
+        # write the modified module back to the file
+        test_path.write_text(modified_module.code, encoding="utf-8")
 
 
 class OptimFunctionCollector(cst.CSTVisitor):

codeflash/optimization/function_optimizer.py

@@ -21,7 +21,11 @@
 from codeflash.benchmarking.utils import process_benchmark_data
 from codeflash.cli_cmds.console import code_print, console, logger, progress_bar
 from codeflash.code_utils import env_utils
-from codeflash.code_utils.code_replacer import replace_function_definitions_in_module
+from codeflash.code_utils.code_replacer import (
+    add_custom_marker_to_all_tests,
+    modify_autouse_fixture,
+    replace_function_definitions_in_module,
+)
 from codeflash.code_utils.code_utils import (
     ImportErrorPattern,
     cleanup_paths,
@@ -742,6 +746,10 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi
                 f"{concolic_coverage_test_files_count} concolic coverage test file"
                 f"{'s' if concolic_coverage_test_files_count != 1 else ''} for {func_qualname}"
             )
+        logger.debug("disabling all autouse fixtures associated with the test files")
+        modify_autouse_fixture(list(unique_instrumented_test_files))
+        logger.debug("add custom marker to all tests")
+        add_custom_marker_to_all_tests(list(unique_instrumented_test_files))
         return unique_instrumented_test_files
 
     def generate_tests_and_optimizations(