start cleaning up and testing

Author: aseembits93

code_to_optimize/tests/pytest/conftest.py

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

code_to_optimize/tests/pytest/test_bubble_sort.py

@@ -16,6 +16,7 @@
 import tomlkit
 from git import InvalidGitRepositoryError, Repo
 from pydantic.dataclasses import dataclass
+from tomlkit import table
 
 from codeflash.api.cfapi import is_github_app_installed_on_repo
 from codeflash.cli_cmds.cli_common import apologize_and_exit, inquirer_wrapper, inquirer_wrapper_path
@@ -728,11 +729,16 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
     tool_section = pyproject_data.get("tool", tomlkit.table())
     tool_section["codeflash"] = codeflash_section
     pyproject_data["tool"] = tool_section
-    if "tool.pytest.ini_options" not in pyproject_data:
-        pyproject_data["tool.pytest.ini_options"] = {}
-    if "markers" not in pyproject_data["tool.pytest.ini_options"]:
-        pyproject_data["tool.pytest.ini_options"]["markers"] = []
-    pyproject_data["tool.pytest.ini_options"]["markers"].append("codeflash_no_autouse")
+    # Create [tool.pytest.ini_options] if it doesn't exist
+    tool_section = pyproject_data.get("tool", table())
+    pytest_section = tool_section.get("pytest", table())
+    ini_options = pytest_section.get("ini_options", table())
+    # Define or overwrite the 'markers' array
+    ini_options["markers"] = ["codeflash_no_autouse"]
+    # Set updated sections back
+    pytest_section["ini_options"] = ini_options
+    tool_section["pytest"] = pytest_section
+    pyproject_data["tool"] = tool_section
     with toml_path.open("w", encoding="utf8") as pyproject_file:
         pyproject_file.write(tomlkit.dumps(pyproject_data))
     click.echo(f"✅ Added Codeflash configuration to {toml_path}")

codeflash/cli_cmds/cmd_init.py

@@ -5,11 +5,9 @@
 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
@@ -23,53 +21,6 @@
     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_extractor.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import ast
-import contextlib
 from collections import defaultdict
 from functools import lru_cache
 from typing import TYPE_CHECKING, Optional, TypeVar
@@ -11,13 +10,9 @@
 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,
-    get_function_qualified_names,
-)
+from codeflash.code_utils.code_extractor import add_global_assignments, add_needed_imports_from_module
 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.code_utils.line_profile_utils import ImportAdder
 from codeflash.models.models import FunctionParent
 
 if TYPE_CHECKING:
@@ -42,6 +37,74 @@ def normalize_code(code: str) -> str:
     return ast.unparse(normalize_node(ast.parse(code)))
 
 
+class PytestMarkAdder(cst.CSTTransformer):
+    """Transformer that adds pytest marks to test functions."""
+
+    def __init__(self, mark_name: str) -> None:
+        super().__init__()
+        self.mark_name = mark_name
+        self.has_pytest_import = False
+
+    def visit_Module(self, node: cst.Module) -> None:
+        """Check if pytest is already imported."""
+        for statement in node.body:
+            if isinstance(statement, cst.SimpleStatementLine):
+                for stmt in statement.body:
+                    if isinstance(stmt, cst.Import):
+                        for import_alias in stmt.names:
+                            if isinstance(import_alias, cst.ImportAlias) and import_alias.name.value == "pytest":
+                                self.has_pytest_import = True
+                    elif isinstance(stmt, cst.ImportFrom) and stmt.module and stmt.module.value == "pytest":
+                        self.has_pytest_import = True
+
+    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:  # noqa: ARG002
+        """Add pytest import if not present."""
+        if not self.has_pytest_import:
+            # Create import statement
+            import_stmt = cst.SimpleStatementLine(body=[cst.Import(names=[cst.ImportAlias(name=cst.Name("pytest"))])])
+            # Add import at the beginning
+            updated_node = updated_node.with_changes(body=[import_stmt, *updated_node.body])
+        return updated_node
+
+    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:  # noqa: ARG002
+        """Add pytest mark to test functions."""
+        # Check if the mark already exists
+        for decorator in updated_node.decorators:
+            if self._is_pytest_mark(decorator.decorator, self.mark_name):
+                return updated_node
+
+        # Create the pytest mark decorator
+        mark_decorator = self._create_pytest_mark()
+
+        # Add the decorator
+        new_decorators = [*list(updated_node.decorators), mark_decorator]
+        return updated_node.with_changes(decorators=new_decorators)
+
+    def _is_pytest_mark(self, decorator: cst.BaseExpression, mark_name: str) -> bool:
+        """Check if a decorator is a specific pytest mark."""
+        if isinstance(decorator, cst.Attribute):
+            if (
+                isinstance(decorator.value, cst.Attribute)
+                and isinstance(decorator.value.value, cst.Name)
+                and decorator.value.value.value == "pytest"
+                and decorator.value.attr.value == "mark"
+                and decorator.attr.value == mark_name
+            ):
+                return True
+        elif isinstance(decorator, cst.Call) and isinstance(decorator.func, cst.Attribute):
+            return self._is_pytest_mark(decorator.func, mark_name)
+        return False
+
+    def _create_pytest_mark(self) -> cst.Decorator:
+        """Create a pytest mark decorator."""
+        # Base: pytest.mark.{mark_name}
+        mark_attr = cst.Attribute(
+            value=cst.Attribute(value=cst.Name("pytest"), attr=cst.Name("mark")), attr=cst.Name(self.mark_name)
+        )
+        decorator = mark_attr
+        return cst.Decorator(decorator=decorator)
+
+
 class AutouseFixtureModifier(cst.CSTTransformer):
     def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
         # Matcher for '@fixture' or '@pytest.fixture'
@@ -63,7 +126,7 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
                 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")')
+                if_test = cst.parse_expression('request.node.get_closest_marker("codeflash_no_autouse")')
                 yield_statement = cst.parse_statement("yield")
                 if_body = cst.IndentedBlock(body=[yield_statement])
 
@@ -75,7 +138,6 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
         return updated_node
 
 
-@contextlib.contextmanager
 def disable_autouse(test_path: Path) -> str:
     file_content = test_path.read_text(encoding="utf-8")
     module = cst.parse_module(file_content)
@@ -107,13 +169,9 @@ def add_custom_marker_to_all_tests(test_paths: list[Path]) -> None:
         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
+        modified_module = cst.parse_module(isort.code(modified_module.code, float_to_top=True))
+        pytest_mark_adder = PytestMarkAdder("codeflash_no_autouse")
+        modified_module = modified_module.visit(pytest_mark_adder)
         test_path.write_text(modified_module.code, encoding="utf-8")
 
 

codeflash/code_utils/code_replacer.py

@@ -213,8 +213,15 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
             for key in set(self.function_to_tests) | set(function_to_concolic_tests)
         }
         instrumented_unittests_created_for_function = self.instrument_existing_tests(function_to_all_tests)
-        # logger.debug("disabling all autouse fixtures associated with the test files")
-        original_conftest_content = modify_autouse_fixture(list(instrumented_unittests_created_for_function))
+        logger.debug("disabling all autouse fixtures associated with the test files")
+        original_conftest_content = modify_autouse_fixture(
+            generated_test_paths + generated_perf_test_paths + list(instrumented_unittests_created_for_function)
+        )
+        logger.debug("add custom marker to all tests")
+        add_custom_marker_to_all_tests(
+            generated_test_paths + generated_perf_test_paths + list(instrumented_unittests_created_for_function)
+        )
+
         # Get a dict of file_path_to_classes of fto and helpers_of_fto
         file_path_to_helper_classes = defaultdict(set)
         for function_source in code_context.helper_functions:
@@ -750,8 +757,6 @@ 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("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(