⚡️ Speed up function _prompt_custom_directory by 363% in PR #1199 (omni-java)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1199

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

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

---

📄 363% (3.63x) speedup for _prompt_custom_directory in codeflash/cli_cmds/init_java.py

⏱️ Runtime : 374 milliseconds → 80.7 milliseconds (best of 34 runs)

📝 Explanation and details

The optimization moved the inquirer.Path question construction out of the while-loop and added @lru_cache(maxsize=1) to _get_theme(), eliminating repeated imports and instantiations of CodeflashTheme on every prompt iteration. The profiler shows _get_theme() was called 1247 times in the original, each time re-importing init_config (~2.2% overhead) and constructing a new theme object (~97.8% overhead, 323 µs per call). Moving the question object outside the loop avoids ~13 µs of reconstruction per iteration, and caching the theme cuts 1246 redundant constructions, yielding a 363% speedup with no functional trade-offs.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import os
from pathlib import Path

import inquirer
import pytest  # used for our unit tests
# import the function under test and the inquirer module so we can monkeypatch its prompt
from codeflash.cli_cmds.init_java import _prompt_custom_directory

def test_returns_string_for_valid_relative_path(monkeypatch):
    # Provide a simple valid relative path as the prompt result.
    # The function should return the string form of whatever was returned under "custom_path".
    monkeypatch.setattr(
        inquirer,
        "prompt",
        lambda questions, theme=None: {"custom_path": "src"},  # first (and only) prompt returns a valid path
    )

    # Call the function and assert the returned value matches the provided path string.
    result = _prompt_custom_directory("project") # 260μs -> 13.7μs (1801% faster)
    assert result == "src"  # returned value must be the string "src"

def test_accepts_pathlib_path_and_returns_str(monkeypatch):
    # Provide a pathlib.Path object as inquirer returns a Path-like object in many backends.
    # The function converts it to string via str(...) before validation/return.
    monkeypatch.setattr(
        inquirer,
        "prompt",
        lambda questions, theme=None: {"custom_path": Path("my_dir/subdir")},
    )

    # The returned value should be the string representation of the Path provided.
    result = _prompt_custom_directory("module") # 263μs -> 20.7μs (1169% faster)
    assert result == "my_dir/subdir"

def test_exits_when_prompt_returns_none(monkeypatch):
    # If inquirer.prompt returns None (user cancelled), the function calls apologize_and_exit(),
    # which results in sys.exit(1) (raised as SystemExit). We assert that behavior.
    monkeypatch.setattr(inquirer, "prompt", lambda questions, theme=None: None)

    with pytest.raises(SystemExit) as excinfo:
        _prompt_custom_directory("config")
    # Confirm it exited with a non-zero status (apologize_and_exit uses sys.exit(1))
    assert excinfo.value.code == 1

def test_invalid_then_valid_loop(monkeypatch):
    # Simulate a first invalid answer (path traversal), then a valid answer.
    responses = [
        {"custom_path": "../etc/passwords"},  # invalid due to '..'
        {"custom_path": "valid_dir"},         # valid relative directory
    ]
    iterator = iter(responses)

    def responder(questions, theme=None):
        # Return the next response on each call to simulate user retrying input.
        return next(iterator)

    monkeypatch.setattr(inquirer, "prompt", responder)

    # Should eventually return the valid path after re-prompting once.
    result = _prompt_custom_directory("data") # 602μs -> 111μs (442% faster)
    assert result == "valid_dir"

@pytest.mark.skipif(os.name == "nt", reason="Null character checks are for unix-like platforms in the validator")
def test_null_character_in_path_is_rejected_then_valid(monkeypatch):
    # On Unix-like systems, a path containing a null character is considered invalid by the validator.
    responses = [
        {"custom_path": "bad\0path"},  # invalid due to null char
        {"custom_path": "okay/path"},  # valid
    ]
    iterator = iter(responses)

    def responder(questions, theme=None):
        return next(iterator)

    monkeypatch.setattr(inquirer, "prompt", responder)

    result = _prompt_custom_directory("assets") # 607μs -> 113μs (434% faster)
    assert result == "okay/path"

@pytest.mark.skipif(os.name != "nt", reason="This absolute path test expects Windows behavior for invalid chars check")

def test_absolute_path_rejected_on_unix_then_valid(monkeypatch, tmp_path):
    # On Unix-like systems, a path starting with '/' is absolute and should be rejected.
    # First response is absolute; second is a valid relative path.
    responses = [
        {"custom_path": str(Path("/etc"))},  # absolute -> invalid on Unix-like platforms
        {"custom_path": "relative_ok"},      # valid
    ]
    iterator = iter(responses)

    def responder(questions, theme=None):
        return next(iterator)

    monkeypatch.setattr(inquirer, "prompt", responder)

    result = _prompt_custom_directory("system") # 607μs -> 116μs (422% faster)
    assert result == "relative_ok"

def test_many_invalid_attempts_then_valid(monkeypatch):
    # Build a large sequence of invalid responses (path traversal attempts) followed by a final valid path.
    # This ensures that the function can handle long re-prompting loops without state corruption.
    invalid_count = 1000  # large-scale number of invalid attempts
    responses = [{"custom_path": f"..{i}/bad"} for i in range(invalid_count)]  # all invalid because they contain '..'
    responses.append({"custom_path": "final_valid_dir"})  # final valid value after many failures

    iterator = iter(responses)

    def responder(questions, theme=None):
        # Return the next response each time inquirer.prompt is called.
        return next(iterator)

    monkeypatch.setattr(inquirer, "prompt", responder)

    # Should eventually return the final valid path after many iterations.
    result = _prompt_custom_directory("large_test") # 303ms -> 67.2ms (352% faster)
    assert result == "final_valid_dir"

def test_mixed_types_in_responses(monkeypatch):
    # Ensure the function handles different types for the "custom_path" value across callbacks:
    # strings and Path objects mixed in sequence.
    responses = [
        {"custom_path": "../bad"},        # invalid
        {"custom_path": Path("ok1")},     # valid Path object
    ]
    iterator = iter(responses)

    def responder(questions, theme=None):
        return next(iterator)

    monkeypatch.setattr(inquirer, "prompt", responder)

    result = _prompt_custom_directory("mixed") # 644μs -> 112μs (472% faster)
    # The function converts the Path to a string before returning, so ensure that behavior.
    assert result == "ok1"

import sys
from pathlib import Path
from unittest.mock import MagicMock, patch

# imports
import pytest
from codeflash.cli_cmds.cli_common import apologize_and_exit
# We need to import the function and its dependencies
from codeflash.cli_cmds.init_java import _prompt_custom_directory
from codeflash.code_utils.code_utils import validate_relative_directory_path

def test_prompt_custom_directory_valid_path_first_try():
    """Test that a valid relative path is returned on the first prompt."""
    # Mock inquirer.prompt to return a valid relative directory path
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "tests"}
        
        # Call the function
        result = _prompt_custom_directory("test") # 273μs -> 27.2μs (906% faster)
        
        # Verify it returns the path as a string
        assert isinstance(result, str)
        assert result == "tests"
        # Verify prompt was called exactly once (valid on first try)
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_valid_nested_path():
    """Test that a valid nested relative path is accepted."""
    # Mock inquirer.prompt to return a nested valid path
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "src/main/java"}
        
        result = _prompt_custom_directory("source") # 270μs -> 28.6μs (844% faster)
        
        # Verify the nested path is returned correctly
        assert result == "src/main/java"
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_invalid_then_valid():
    """Test that invalid paths are rejected and valid paths are accepted on retry."""
    # Mock inquirer.prompt to return invalid path first, then valid path
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.click.echo') as mock_echo:
            # First call returns invalid (absolute path), second call returns valid
            mock_prompt.side_effect = [
                {"custom_path": "/absolute/path"},  # Invalid - absolute path
                {"custom_path": "valid_relative"},  # Valid
            ]
            
            result = _prompt_custom_directory("config")
            
            # Verify the valid path is returned
            assert result == "valid_relative"
            # Verify prompt was called twice (once for invalid, once for valid)
            assert mock_prompt.call_count == 2
            # Verify error messages were displayed for the invalid path
            assert mock_echo.call_count >= 2

def test_prompt_custom_directory_user_cancels():
    """Test that the function exits when user cancels (Ctrl+C)."""
    # Mock inquirer.prompt to return None (user cancelled)
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.apologize_and_exit') as mock_exit:
            mock_prompt.return_value = None
            mock_exit.side_effect = SystemExit(1)
            
            # Call should trigger exit
            with pytest.raises(SystemExit):
                _prompt_custom_directory("test")
            
            # Verify apologize_and_exit was called
            mock_exit.assert_called_once()

def test_prompt_custom_directory_message_includes_dir_type():
    """Test that the prompt message includes the directory type parameter."""
    # Mock inquirer.prompt and click.echo to capture messages
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "src"}
        
        _prompt_custom_directory("source") # 266μs -> 26.0μs (925% faster)
        
        # Verify that inquirer.Path was created with the correct message
        # Check the call to inquirer.prompt includes the expected questions
        assert mock_prompt.call_count == 1
        # The call_args[0][0] should be the list of questions
        questions = mock_prompt.call_args[0][0]
        assert len(questions) > 0
        # Check that the message contains the dir_type
        assert "source" in questions[0].message

def test_prompt_custom_directory_empty_string_path():
    """Test that empty string paths are rejected."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.click.echo') as mock_echo:
            # First empty path, then valid path
            mock_prompt.side_effect = [
                {"custom_path": ""},  # Invalid - empty
                {"custom_path": "valid"},  # Valid
            ]
            
            result = _prompt_custom_directory("test")
            
            assert result == "valid"
            assert mock_prompt.call_count == 2

def test_prompt_custom_directory_whitespace_only_path():
    """Test that whitespace-only paths are rejected."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.click.echo') as mock_echo:
            # First whitespace-only, then valid
            mock_prompt.side_effect = [
                {"custom_path": "   "},  # Invalid - whitespace only
                {"custom_path": "valid"},  # Valid
            ]
            
            result = _prompt_custom_directory("test")
            
            assert result == "valid"
            assert mock_prompt.call_count == 2

def test_prompt_custom_directory_path_traversal_attempt():
    """Test that path traversal attempts (..) are rejected."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.click.echo') as mock_echo:
            # First with path traversal, then valid
            mock_prompt.side_effect = [
                {"custom_path": "../../../etc/passwd"},  # Invalid - traversal
                {"custom_path": "src"},  # Valid
            ]
            
            result = _prompt_custom_directory("test")
            
            assert result == "src"
            assert mock_prompt.call_count == 2

def test_prompt_custom_directory_single_character_path():
    """Test that a single character valid path is accepted."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "a"}
        
        result = _prompt_custom_directory("test") # 265μs -> 25.5μs (942% faster)
        
        assert result == "a"
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_path_with_dots():
    """Test that valid paths with dots (but not ..) are accepted."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "src.backup"}
        
        result = _prompt_custom_directory("test") # 268μs -> 25.2μs (965% faster)
        
        assert result == "src.backup"
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_path_with_underscores():
    """Test that paths with underscores are accepted."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "test_src_main"}
        
        result = _prompt_custom_directory("test") # 266μs -> 25.6μs (941% faster)
        
        assert result == "test_src_main"
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_path_with_hyphens():
    """Test that paths with hyphens are accepted."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "my-src-main"}
        
        result = _prompt_custom_directory("test") # 267μs -> 25.6μs (946% faster)
        
        assert result == "my-src-main"
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_different_dir_types():
    """Test prompting for various directory types."""
    dir_types = ["test", "source", "config", "build", "artifact"]
    
    for dir_type in dir_types:
        with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
            mock_prompt.return_value = {"custom_path": "src"}
            
            result = _prompt_custom_directory(dir_type)
            
            # Verify result is consistent
            assert result == "src"

def test_prompt_custom_directory_path_with_leading_dot():
    """Test that paths starting with dot (current directory) are handled."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        # Note: "./tests" is relative and should be valid
        mock_prompt.return_value = {"custom_path": "./tests"}
        
        result = _prompt_custom_directory("test") # 265μs -> 26.3μs (912% faster)
        
        # The function should accept this as it's a valid relative path
        assert result == "./tests"
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_multiple_invalid_attempts():
    """Test that multiple invalid attempts eventually lead to a valid result."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.click.echo') as mock_echo:
            # Multiple invalid attempts followed by valid
            mock_prompt.side_effect = [
                {"custom_path": "/absolute"},  # Invalid 1
                {"custom_path": "../relative"},  # Invalid 2
                {"custom_path": ""},  # Invalid 3
                {"custom_path": "   "},  # Invalid 4
                {"custom_path": "finally_valid"},  # Valid
            ]
            
            result = _prompt_custom_directory("test")
            
            assert result == "finally_valid"
            assert mock_prompt.call_count == 5

def test_prompt_custom_directory_case_sensitivity():
    """Test that path case is preserved."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        mock_prompt.return_value = {"custom_path": "MyTestDir"}
        
        result = _prompt_custom_directory("test") # 266μs -> 25.6μs (940% faster)
        
        # Case should be preserved in the output
        assert result == "MyTestDir"

def test_prompt_custom_directory_unicode_path():
    """Test behavior with unicode characters in path."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        # Some unicode might be invalid on some systems, so we test valid unicode
        mock_prompt.return_value = {"custom_path": "src_café"}
        
        result = _prompt_custom_directory("test") # 267μs -> 25.3μs (959% faster)
        
        # Should accept unicode if it's valid on the platform
        assert isinstance(result, str)

def test_prompt_custom_directory_long_path():
    """Test that appropriately long paths are handled."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        long_path = "a/b/c/d/e/f/g/h/i/j/k/l/m/n/o/p/q/r/s/t/u/v/w/x/y/z"
        mock_prompt.return_value = {"custom_path": long_path}
        
        result = _prompt_custom_directory("test") # 279μs -> 32.7μs (756% faster)
        
        assert result == long_path

def test_prompt_custom_directory_many_retries_still_works():
    """Test that the function handles many retry attempts without breaking."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.click.echo') as mock_echo:
            # Generate many invalid paths followed by one valid
            invalid_paths = [{"custom_path": f"/invalid_{i}"} for i in range(100)]
            invalid_paths.append({"custom_path": "valid_finally"})
            
            mock_prompt.side_effect = invalid_paths
            
            result = _prompt_custom_directory("test")
            
            assert result == "valid_finally"
            assert mock_prompt.call_count == 101

def test_prompt_custom_directory_very_deep_nested_path():
    """Test with a very deeply nested relative path."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        # Create a deeply nested path
        deep_path = "/".join([f"level{i}" for i in range(100)])
        mock_prompt.return_value = {"custom_path": deep_path}
        
        result = _prompt_custom_directory("test") # 348μs -> 63.4μs (450% faster)
        
        assert result == deep_path
        assert mock_prompt.call_count == 1

def test_prompt_custom_directory_path_with_many_segments():
    """Test path with many directory segments."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        # Create path with 500 segments
        path_segments = [f"dir{i}" for i in range(500)]
        complex_path = "/".join(path_segments)
        mock_prompt.return_value = {"custom_path": complex_path}
        
        result = _prompt_custom_directory("test") # 574μs -> 180μs (219% faster)
        
        assert result == complex_path

def test_prompt_custom_directory_alternating_valid_invalid():
    """Test alternating between valid and invalid paths in retries."""
    with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
        with patch('codeflash.cli_cmds.init_java.click.echo') as mock_echo:
            # Alternate between invalid and valid 50 times, then end with valid
            paths = []
            for i in range(50):
                paths.append({"custom_path": f"/invalid_{i}"})  # Invalid
                paths.append({"custom_path": f"valid_{i}"})  # Valid but still invalid due to next invalid
            paths.append({"custom_path": "final_valid"})
            
            mock_prompt.side_effect = paths
            
            result = _prompt_custom_directory("test")
            
            assert result == "final_valid"

def test_prompt_custom_directory_stress_test_return_type():
    """Stress test: verify return type consistency across many calls."""
    for i in range(100):
        with patch('codeflash.cli_cmds.init_java.inquirer.prompt') as mock_prompt:
            mock_prompt.return_value = {"custom_path": f"test_path_{i}"}
            
            result = _prompt_custom_directory("test")
            
            # Every result must be a string
            assert isinstance(result, str)
            assert result == f"test_path_{i}"

To edit these changes git checkout codeflash/optimize-pr1199-2026-03-13T01.44.29 and push.

[claude]

Claude finished @codeflash-ai[bot]'s task in 3m 4s —— View job

---

PR Review Summary

Todo

• Triage PR scope
• Lint and typecheck
• Resolve stale threads
• Code review
• Duplicate detection
• Test coverage
• Check optimization PRs

---

Prek Checks

✅ All checks pass (ruff check, ruff format). No issues introduced by this PR.

Mypy

Three pre-existing errors in init_java.py — all present in origin/omni-java before this PR:

• _get_theme() missing return type annotation (line 60) — pre-existing
• Returning Any in _prompt_custom_directory (line 366) — pre-existing
• Unsupported target for indexed assignment (line 496) — pre-existing

No new type errors introduced.

Code Review

Correctness of inquirer.Path object reuse (line 371–377):

The Question base class stores mutable state in self.answers = {}. However, inspecting ConsoleRender.render() shows that question.answers = answers or {} is called at the start of every prompt() invocation — resetting any prior state. Reusing the custom_question object across loop iterations is safe.

Correctness of @lru_cache(maxsize=1) on _get_theme() (line 59):

_get_theme() takes no arguments, so the first CodeflashTheme() instance is cached for the process lifetime. If CodeflashTheme accumulates mutable state during use (e.g., via a render call), all callers of _get_theme() will share that state. This is worth confirming — but themes in inquirer are typically read-only configuration objects, so this appears safe.

No bugs found.

Duplicate Detection

HIGH confidence duplicate: init_javascript.py contains an identical _get_theme() function (line 74) and an identical pre-optimization _prompt_custom_directory() function (line 480) that still constructs the inquirer.Path question inside the loop and calls _get_theme() without caching.

The same optimization applied here could be applied to init_javascript.py for consistency. The two functions are semantically identical. Fix this →%20Add%20%40lru_cache(maxsize%3D1)%20to%20_get_theme()%20at%20line%2074%2C%20and%202)%20move%20the%20inquirer.Path%20question%20construction%20outside%20the%20while%20loop%20in%20_prompt_custom_directory()%20at%20line%20480%20on%20branch%20codeflash%2Foptimize-pr1199-2026-03-13T01.44.29&repo=codeflash-ai/codeflash)

Test Coverage

No existing unit tests for codeflash/cli_cmds/init_java.py (not new — pre-existing gap). The PR's 134 generated regression tests verified correctness but are not committed to the repository.

Optimization PRs

Three optimization PRs are currently open against omni-java, all from today (CI still pending):

• ⚡️ Speed up function _prompt_custom_directory by 363% in PR #1199 (omni-java) #1827 (this PR) — _prompt_custom_directory 3.63x speedup
• ⚡️ Speed up method TestGenRequest.to_payload by 20% in PR #1199 (omni-java) #1826 — TestGenRequest.to_payload 1.20x speedup
• ⚡️ Speed up method OptimizeRequest.to_payload by 33% in PR #1199 (omni-java) #1825 — OptimizeRequest.to_payload 1.33x speedup

CI is still running on all three — will not merge until checks complete.

---

Overall: The optimization is correct and the speedup is legitimate. The only notable finding is the duplicate unoptimized code in init_javascript.py.