⚡️ Speed up function _handle_reset_config by 13% in PR #1723 (omni-main-java)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1723

If you approve this dependent PR, these changes will be merged into the original PR branch omni-main-java.

This PR will be automatically closed if the original PR is merged.

---

📄 13% (0.13x) speedup for _handle_reset_config in codeflash/cli_cmds/cli.py

⏱️ Runtime : 16.3 milliseconds → 14.4 milliseconds (best of 5 runs)

📝 Explanation and details

The hot path in has_existing_config spent 93% of its runtime parsing TOML files with tomlkit.parse() and 0.4% parsing JSON with json.load(), both invoked unconditionally on every existing file. The optimization replaces binary mode with text mode ("r", encoding="utf8", errors="ignore") and performs a quick substring search ("codeflash" in content) before attempting to parse, exiting early if the target key isn't present. This reduces average time per call from ~7.4 ms to ~1.2 ms (6.2×), which cascades to shave 18.9% off _handle_reset_config and yields a 13% end-to-end runtime improvement. No correctness regressions: all tests pass with identical behavior.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 45 Passed
🌀 Generated Regression Tests	✅ 8 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	96.4%

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_setup/test_e2e_setup.py::TestE2ECLIFlags.test_reset_config_handles_no_config	1.32ms	278μs	376%✅
test_setup/test_e2e_setup.py::TestE2ECLIFlags.test_reset_config_removes_from_package_json	726μs	703μs	3.30%✅
test_setup/test_e2e_setup.py::TestE2ECLIFlags.test_reset_config_removes_from_pyproject	3.47ms	3.00ms	15.6%✅

🌀 Click to see Generated Regression Tests

import json
import os
from pathlib import Path

import pytest
import tomlkit
from codeflash.cli_cmds import \
    console as console_module  # real console module used by the CLI
from codeflash.cli_cmds.cli import _handle_reset_config
from codeflash.cli_cmds.console import console

# Helper fake console class used in tests to capture prints and control input behavior.
class FakeConsole:
    def __init__(self, inputs=None, raise_on_input=None):
        # inputs: an iterable or callable providing responses for input()
        # raise_on_input: if set to an Exception instance, input() will raise it
        self.captured = []  # store printed messages here
        self._inputs = list(inputs) if inputs is not None and not callable(inputs) else inputs
        self._raise = raise_on_input

    def print(self, *args, **kwargs):
        # Convert all args to strings and capture a single joined message (like console)
        self.captured.append(" ".join(str(a) for a in args))

    def input(self, prompt=""):
        # Capture the prompt too so tests can assert it was shown
        self.captured.append(f"__INPUT_PROMPT__ {prompt}")
        if self._raise:
            # Simulate user abort (EOF or KeyboardInterrupt)
            raise self._raise
        if callable(self._inputs):
            # If a callable is provided, call it to get a response
            return self._inputs()
        if self._inputs:
            return self._inputs.pop(0)
        # Default to empty string (treated as "no")
        return ""

# Utility context manager to temporarily change working directory
class TempCwd:
    def __init__(self, new_cwd: Path):
        self.new_cwd = new_cwd
        self.old_cwd = Path.cwd()

    def __enter__(self):
        os.chdir(self.new_cwd)
        return self

    def __exit__(self, exc_type, exc, tb):
        os.chdir(self.old_cwd)

def test_no_existing_config_prints_message(tmp_path: Path):
    # Create an empty project directory (no config files).
    project_dir = tmp_path / "proj_empty"
    project_dir.mkdir()

    fake = FakeConsole()
    orig_console = console_module.console
    try:
        # Replace the real console with our fake console so we can capture output.
        console_module.console = fake

        # Run in the temporary directory where there is no config.
        with TempCwd(project_dir):
            # confirm=False avoids prompting; function should detect no existing config and print a message.
            _handle_reset_config(confirm=False)

        # There should be at least one captured message indicating nothing to remove.
        joined = "\n".join(fake.captured)
    finally:
        # Restore original console to avoid side effects.
        console_module.console = orig_console

def test_remove_from_pyproject_confirm_false(tmp_path: Path):
    # Create a project with a pyproject.toml containing [tool.codeflash]
    project_dir = tmp_path / "proj_py"
    project_dir.mkdir()
    pyproject = project_dir / "pyproject.toml"
    # Write valid TOML with [tool.codeflash]
    doc = tomlkit.document()
    doc["tool"] = {"codeflash": {"some_option": True}}
    pyproject.write_text(tomlkit.dumps(doc), encoding="utf8")

    fake = FakeConsole()
    orig_console = console_module.console
    try:
        console_module.console = fake

        with TempCwd(project_dir):
            # confirm=False so no interactive prompt is needed; should remove the section.
            _handle_reset_config(confirm=False)

        # After removal, pyproject.toml should no longer include "codeflash" under [tool]
        content = pyproject.read_text(encoding="utf8")

        # The fake console should have printed a success message with a green check mark
        # and the removal message (brackets escaped as "\[").
        joined = "\n".join(fake.captured)
    finally:
        console_module.console = orig_console

def test_remove_from_package_json_confirm_false(tmp_path: Path):
    # Create a project with package.json containing "codeflash" key.
    project_dir = tmp_path / "proj_pkg"
    project_dir.mkdir()
    package_json = project_dir / "package.json"
    data = {"name": "proj", "codeflash": {"opt": 1}}
    package_json.write_text(json.dumps(data, indent=2) + "\n", encoding="utf8")

    fake = FakeConsole()
    orig_console = console_module.console
    try:
        console_module.console = fake

        with TempCwd(project_dir):
            # confirm=False so it will remove from package.json
            _handle_reset_config(confirm=False)

        # package.json should no longer contain the "codeflash" key
        content = json.loads(package_json.read_text(encoding="utf8"))

        # Confirm success printed
        joined = "\n".join(fake.captured)
    finally:
        console_module.console = orig_console

def test_confirm_prompt_user_denies(tmp_path: Path):
    # When user denies the prompt, no removal should occur and "Cancelled." is printed.
    project_dir = tmp_path / "proj_confirm_no"
    project_dir.mkdir()
    pyproject = project_dir / "pyproject.toml"
    doc = tomlkit.document()
    doc["tool"] = {"codeflash": {"flag": True}}
    pyproject.write_text(tomlkit.dumps(doc), encoding="utf8")

    # Fake console returns "n" for the prompt (deny)
    fake = FakeConsole(inputs=["n"])
    orig_console = console_module.console
    try:
        console_module.console = fake

        with TempCwd(project_dir):
            # confirm=True so the function will prompt.
            _handle_reset_config(confirm=True)

        # File should remain unchanged (still contains codeflash)
        content = pyproject.read_text(encoding="utf8")

        # Should have printed a cancellation message
        joined = "\n".join(fake.captured)
    finally:
        console_module.console = orig_console

def test_confirm_prompt_user_accepts(tmp_path: Path):
    # When user accepts (yes), the config should be removed.
    project_dir = tmp_path / "proj_confirm_yes"
    project_dir.mkdir()
    package_json = project_dir / "package.json"
    data = {"name": "x", "codeflash": {"a": 1}}
    package_json.write_text(json.dumps(data, indent=2) + "\n", encoding="utf8")

    # Fake console returns "yes" for the prompt
    fake = FakeConsole(inputs=["yes"])
    orig_console = console_module.console
    try:
        console_module.console = fake

        with TempCwd(project_dir):
            _handle_reset_config(confirm=True)

        # Confirm the package.json no longer has the codeflash key
        loaded = json.loads(package_json.read_text(encoding="utf8"))

        # Should have printed success indicator
        joined = "\n".join(fake.captured)
    finally:
        console_module.console = orig_console

def test_input_raises_eoferror_cancels(tmp_path: Path):
    # Simulate an EOFError during input; should print cancelled and leave file intact.
    project_dir = tmp_path / "proj_eof"
    project_dir.mkdir()
    pyproject = project_dir / "pyproject.toml"
    doc = tomlkit.document()
    doc["tool"] = {"codeflash": {"k": "v"}}
    pyproject.write_text(tomlkit.dumps(doc), encoding="utf8")

    # Fake console that raises EOFError when input() is called
    fake = FakeConsole(raise_on_input=EOFError())
    orig_console = console_module.console
    try:
        console_module.console = fake

        with TempCwd(project_dir):
            _handle_reset_config(confirm=True)

        # File should remain intact
        content = pyproject.read_text(encoding="utf8")

        # Should have printed cancellation message
        joined = "\n".join(fake.captured)
    finally:
        console_module.console = orig_console

def test_invalid_pyproject_failure_message(tmp_path: Path):
    # If pyproject is present but invalid such that removal fails, the CLI should print an error (red X).
    project_dir = tmp_path / "proj_invalid"
    project_dir.mkdir()
    pyproject = project_dir / "pyproject.toml"
    # Write invalid TOML so tomlkit.parse will raise and removal returns False
    pyproject.write_bytes(b"\xff\xff not a toml")

    fake = FakeConsole()
    orig_console = console_module.console
    try:
        console_module.console = fake

        with TempCwd(project_dir):
            _handle_reset_config(confirm=False)

        # The output should contain a failure indicator (red X) and "Failed to remove config" text.
        joined = "\n".join(fake.captured)
    finally:
        console_module.console = orig_console

def test_large_scale_package_json_removal_with_many_js_files(tmp_path: Path):
    # Create a project with one package.json containing codeflash config and many .js files
    project_dir = tmp_path / "proj_big"
    project_dir.mkdir()
    package_json = project_dir / "package.json"
    data = {"name": "big", "codeflash": {"enabled": True}}
    package_json.write_text(json.dumps(data, indent=2) + "\n", encoding="utf8")

    # Create a large number of .js files (1000) nested across several directories to exercise rglob scanning.
    count = 1000
    for i in range(count):
        # Spread files across 20 subdirectories to avoid filesystem limits with too many files in one dir.
        sub = project_dir / f"subdir_{i % 20}"
        sub.mkdir(exist_ok=True)
        f = sub / f"file_{i}.test.js"
        f.write_text(f"// JS test file {i}\n", encoding="utf8")

    fake = FakeConsole()
    orig_console = console_module.console
    try:
        console_module.console = fake

        with TempCwd(project_dir):
            # confirm=False to directly remove package.json codeflash section
            _handle_reset_config(confirm=False)

        # Ensure package.json no longer contains "codeflash" key
        loaded = json.loads(package_json.read_text(encoding="utf8"))

        # Check that operation reported success
        joined = "\n".join(fake.captured)
    finally:
        console_module.console = orig_console
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1723-2026-03-02T21.55.08 and push.