follow up

Author: KRRT7

codeflash/cli_cmds/cmd_init.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import ast
 import os
 import re
 import subprocess
@@ -59,7 +58,6 @@ class CLISetupInfo:
     module_root: str
     tests_root: str
     benchmarks_root: Union[str, None]
-    test_framework: str
     ignore_paths: list[str]
     formatter: Union[str, list[str]]
     git_remote: str
@@ -70,7 +68,6 @@ class CLISetupInfo:
 class VsCodeSetupInfo:
     module_root: str
     tests_root: str
-    test_framework: str
     formatter: Union[str, list[str]]
 
 
@@ -256,7 +253,6 @@ def __init__(self) -> None:
 class CommonSections(Enum):
     module_root = "module_root"
     tests_root = "tests_root"
-    test_framework = "test_framework"
     formatter_cmds = "formatter_cmds"
 
     def get_toml_key(self) -> str:
@@ -292,9 +288,6 @@ def get_suggestions(section: str) -> tuple[list[str], Optional[str]]:
     if section == CommonSections.tests_root:
         default = "tests" if "tests" in valid_subdirs else None
         return valid_subdirs, default
-    if section == CommonSections.test_framework:
-        auto_detected = detect_test_framework_from_config_files(Path.cwd())
-        return ["pytest", "unittest"], auto_detected
     if section == CommonSections.formatter_cmds:
         return ["disabled", "ruff", "black"], "disabled"
     msg = f"Unknown section: {section}"
@@ -455,43 +448,6 @@ def collect_setup_info() -> CLISetupInfo:
 
     ph("cli-tests-root-provided")
 
-    test_framework_choices, detected_framework = get_suggestions(CommonSections.test_framework)
-    autodetected_test_framework = detected_framework or detect_test_framework_from_test_files(tests_root)
-
-    framework_message = "⚗️ Let's configure your test framework.\n\n"
-    if autodetected_test_framework:
-        framework_message += f"I detected that you're using {autodetected_test_framework}. "
-    framework_message += "Please confirm or select a different one."
-
-    framework_panel = Panel(Text(framework_message, style="blue"), title="⚗️ Test Framework", border_style="bright_blue")
-    console.print(framework_panel)
-    console.print()
-
-    framework_choices = []
-    # add icons based on the detected framework
-    for choice in test_framework_choices:
-        if choice == "pytest":
-            framework_choices.append(("🧪 pytest", "pytest"))
-        elif choice == "unittest":
-            framework_choices.append(("🐍 unittest", "unittest"))
-
-    framework_questions = [
-        inquirer.List(
-            "test_framework",
-            message="Which test framework do you use?",
-            choices=framework_choices,
-            default=autodetected_test_framework or "pytest",
-            carousel=True,
-        )
-    ]
-
-    framework_answers = inquirer.prompt(framework_questions, theme=CodeflashTheme())
-    if not framework_answers:
-        apologize_and_exit()
-    test_framework = framework_answers["test_framework"]
-
-    ph("cli-test-framework-provided", {"test_framework": test_framework})
-
     benchmarks_root = None
 
     # TODO: Implement other benchmark framework options
@@ -588,60 +544,13 @@ def collect_setup_info() -> CLISetupInfo:
         module_root=str(module_root),
         tests_root=str(tests_root),
         benchmarks_root=str(benchmarks_root) if benchmarks_root else None,
-        test_framework=cast("str", test_framework),
         ignore_paths=ignore_paths,
         formatter=cast("str", formatter),
         git_remote=str(git_remote),
         enable_telemetry=enable_telemetry,
     )
 
 
-def detect_test_framework_from_config_files(curdir: Path) -> Optional[str]:
-    test_framework = None
-    pytest_files = ["pytest.ini", "pyproject.toml", "tox.ini", "setup.cfg"]
-    pytest_config_patterns = {
-        "pytest.ini": "[pytest]",
-        "pyproject.toml": "[tool.pytest.ini_options]",
-        "tox.ini": "[pytest]",
-        "setup.cfg": "[tool:pytest]",
-    }
-    for pytest_file in pytest_files:
-        file_path = curdir / pytest_file
-        if file_path.exists():
-            with file_path.open(encoding="utf8") as file:
-                contents = file.read()
-                if pytest_config_patterns[pytest_file] in contents:
-                    test_framework = "pytest"
-                    break
-        test_framework = "pytest"
-    return test_framework
-
-
-def detect_test_framework_from_test_files(tests_root: Path) -> Optional[str]:
-    test_framework = None
-    # Check if any python files contain a class that inherits from unittest.TestCase
-    for filename in tests_root.iterdir():
-        if filename.suffix == ".py":
-            with filename.open(encoding="utf8") as file:
-                contents = file.read()
-                try:
-                    node = ast.parse(contents)
-                except SyntaxError:
-                    continue
-                if any(
-                    isinstance(item, ast.ClassDef)
-                    and any(
-                        (isinstance(base, ast.Attribute) and base.attr == "TestCase")
-                        or (isinstance(base, ast.Name) and base.id == "TestCase")
-                        for base in item.bases
-                    )
-                    for item in node.body
-                ):
-                    test_framework = "unittest"
-                    break
-    return test_framework
-
-
 def check_for_toml_or_setup_file() -> str | None:
     click.echo()
     click.echo("Checking for pyproject.toml or setup.py…\r", nl=False)
@@ -1060,7 +969,6 @@ def configure_pyproject_toml(
     else:
         codeflash_section["module-root"] = setup_info.module_root
         codeflash_section["tests-root"] = setup_info.tests_root
-        codeflash_section["test-framework"] = setup_info.test_framework
         codeflash_section["ignore-paths"] = setup_info.ignore_paths
         if not setup_info.enable_telemetry:
             codeflash_section["disable-telemetry"] = not setup_info.enable_telemetry

codeflash/code_utils/config_parser.py

@@ -149,10 +149,6 @@ def parse_config_file(
         else:
             config[key] = []
 
-    if config.get("test-framework"):
-        assert config["test-framework"] in {"pytest", "unittest"}, (
-            "In pyproject.toml, Codeflash only supports the 'test-framework' as pytest and unittest."
-        )
     # see if this is happening during GitHub actions setup
     if config.get("formatter-cmds") and len(config.get("formatter-cmds")) > 0 and not override_formatter_check:
         assert config.get("formatter-cmds")[0] != "your-formatter $file", (

codeflash/lsp/beta.py

@@ -187,7 +187,6 @@ def write_config(params: WriteConfigParams) -> dict[str, any]:
     setup_info = VsCodeSetupInfo(
         module_root=getattr(cfg, "module_root", ""),
         tests_root=getattr(cfg, "tests_root", ""),
-        test_framework=getattr(cfg, "test_framework", "pytest"),
         formatter=get_formatter_cmds(getattr(cfg, "formatter_cmds", "disabled")),
     )
 
@@ -203,7 +202,6 @@ def write_config(params: WriteConfigParams) -> dict[str, any]:
 def get_config_suggestions(_params: any) -> dict[str, any]:
     module_root_suggestions, default_module_root = get_suggestions(CommonSections.module_root)
     tests_root_suggestions, default_tests_root = get_suggestions(CommonSections.tests_root)
-    test_framework_suggestions, default_test_framework = get_suggestions(CommonSections.test_framework)
     formatter_suggestions, default_formatter = get_suggestions(CommonSections.formatter_cmds)
     get_valid_subdirs.cache_clear()
 
@@ -238,7 +236,6 @@ def get_config_suggestions(_params: any) -> dict[str, any]:
     return {
         "module_root": {"choices": module_root_suggestions, "default": default_module_root},
         "tests_root": {"choices": tests_root_suggestions, "default": default_tests_root},
-        "test_framework": {"choices": test_framework_suggestions, "default": default_test_framework},
         "formatter_cmds": {"choices": formatter_suggestions, "default": default_formatter},
     }
 

pyproject.toml

@@ -303,7 +303,6 @@ __version__ = "{version}"
 module-root = "codeflash"
 tests-root = "tests"
 benchmarks-root = "tests/benchmarks"
-test-framework = "pytest"
 formatter-cmds = [
     "uvx ruff check --exit-zero --fix $file",
     "uvx ruff format $file",

tests/test_cmd_init.py

@@ -1,15 +1,17 @@
-import pytest
+import os
 import tempfile
 from pathlib import Path
+
+import pytest
+
 from codeflash.cli_cmds.cmd_init import (
-    is_valid_pyproject_toml,
-    configure_pyproject_toml,
     CLISetupInfo,
-    get_formatter_cmds,
     VsCodeSetupInfo,
+    configure_pyproject_toml,
+    get_formatter_cmds,
     get_valid_subdirs,
+    is_valid_pyproject_toml,
 )
-import os
 
 
 @pytest.fixture
@@ -29,11 +31,12 @@ def test_is_valid_pyproject_toml_with_empty_config(temp_dir: Path) -> None:
         assert not valid
         assert _message == "Missing required field: 'module_root'"
 
+
 def test_is_valid_pyproject_toml_with_incorrect_module_root(temp_dir: Path) -> None:
     with (temp_dir / "pyproject.toml").open(mode="w") as f:
         wrong_module_root = temp_dir / "invalid_directory"
         f.write(
-            f"""[tool.codeflash]
+            """[tool.codeflash]
 module-root = "invalid_directory"
 """
         )
@@ -47,7 +50,7 @@ def test_is_valid_pyproject_toml_with_incorrect_tests_root(temp_dir: Path) -> No
     with (temp_dir / "pyproject.toml").open(mode="w") as f:
         wrong_tests_root = temp_dir / "incorrect_tests_root"
         f.write(
-            f"""[tool.codeflash]
+            """[tool.codeflash]
 module-root = "."
 tests-root = "incorrect_tests_root"
 """
@@ -65,21 +68,21 @@ def test_is_valid_pyproject_toml_with_valid_config(temp_dir: Path) -> None:
             """[tool.codeflash]
 module-root = "."
 tests-root = "tests"
-test-framework = "pytest"
 """
         )
         f.flush()
         valid, config, _message = is_valid_pyproject_toml(temp_dir / "pyproject.toml")
         assert valid
 
+
 def test_get_formatter_cmd(temp_dir: Path) -> None:
     assert get_formatter_cmds("black") == ["black $file"]
     assert get_formatter_cmds("ruff") == ["ruff check --exit-zero --fix $file", "ruff format $file"]
     assert get_formatter_cmds("disabled") == ["disabled"]
     assert get_formatter_cmds("don't use a formatter") == ["disabled"]
 
-def test_configure_pyproject_toml_for_cli(temp_dir: Path) -> None:
 
+def test_configure_pyproject_toml_for_cli(temp_dir: Path) -> None:
     pyproject_path = temp_dir / "pyproject.toml"
 
     with (pyproject_path).open(mode="w") as f:
@@ -90,7 +93,6 @@ def test_configure_pyproject_toml_for_cli(temp_dir: Path) -> None:
             module_root=".",
             tests_root="tests",
             benchmarks_root=None,
-            test_framework="pytest",
             ignore_paths=[],
             formatter="black",
             git_remote="origin",
@@ -101,81 +103,78 @@ def test_configure_pyproject_toml_for_cli(temp_dir: Path) -> None:
         assert success
 
         config_content = pyproject_path.read_text()
-        assert """[tool.codeflash]
+        assert (
+            config_content
+            == """[tool.codeflash]
 # All paths are relative to this pyproject.toml's directory.
 module-root = "."
 tests-root = "tests"
-test-framework = "pytest"
 ignore-paths = []
 disable-telemetry = true
 formatter-cmds = ["black $file"]
-""" == config_content
+"""
+        )
         valid, _, _ = is_valid_pyproject_toml(pyproject_path)
         assert valid
 
-def test_configure_pyproject_toml_for_vscode_with_empty_config(temp_dir: Path) -> None:
 
+def test_configure_pyproject_toml_for_vscode_with_empty_config(temp_dir: Path) -> None:
     pyproject_path = temp_dir / "pyproject.toml"
 
     with (pyproject_path).open(mode="w") as f:
         f.write("")
         f.flush()
         os.mkdir(temp_dir / "tests")
-        config = VsCodeSetupInfo(
-            module_root=".",
-            tests_root="tests",
-            test_framework="pytest",
-            formatter="black",
-        )
+        config = VsCodeSetupInfo(module_root=".", tests_root="tests", formatter="black")
 
         success = configure_pyproject_toml(config, pyproject_path)
         assert success
 
         config_content = pyproject_path.read_text()
-        assert """[tool.codeflash]
+        assert (
+            config_content
+            == """[tool.codeflash]
 module-root = "."
 tests-root = "tests"
-test-framework = "pytest"
 formatter-cmds = ["black $file"]
-""" == config_content
+"""
+        )
         valid, _, _ = is_valid_pyproject_toml(pyproject_path)
         assert valid
 
+
 def test_configure_pyproject_toml_for_vscode_with_existing_config(temp_dir: Path) -> None:
     pyproject_path = temp_dir / "pyproject.toml"
-    
+
     with (pyproject_path).open(mode="w") as f:
         f.write("""[tool.codeflash]
 module-root = "codeflash"
 tests-root = "tests"
 benchmarks-root = "tests/benchmarks"
-test-framework = "pytest"
 formatter-cmds = ["disabled"]
 """)
         f.flush()
         os.mkdir(temp_dir / "tests")
-        config = VsCodeSetupInfo(
-            module_root=".",
-            tests_root="tests",
-            test_framework="pytest",
-            formatter="disabled",
-        )
+        config = VsCodeSetupInfo(module_root=".", tests_root="tests", formatter="disabled")
 
         success = configure_pyproject_toml(config, pyproject_path)
         assert success
 
         config_content = pyproject_path.read_text()
         # the benchmarks-root shouldn't get overwritten
-        assert """[tool.codeflash]
+        assert (
+            config_content
+            == """[tool.codeflash]
 module-root = "."
 tests-root = "tests"
 benchmarks-root = "tests/benchmarks"
-test-framework = "pytest"
 formatter-cmds = ["disabled"]
-""" == config_content
+"""
+        )
         valid, _, _ = is_valid_pyproject_toml(pyproject_path)
         assert valid
 
+
 def test_get_valid_subdirs(temp_dir: Path) -> None:
     os.mkdir(temp_dir / "dir1")
     os.mkdir(temp_dir / "dir2")