Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions codeflash/api/aiservice.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ def optimize_python_code_line_profiler( # noqa: D417
line_profiler_results: str,
n_candidates: int,
experiment_metadata: ExperimentMetadata | None = None,
is_numerical_code: bool | None = None, # noqa: FBT001
is_numerical_code: bool | None = None,
) -> list[OptimizedCandidate]:
"""Optimize the given python code for performance using line profiler results.

Expand Down Expand Up @@ -646,7 +646,7 @@ def generate_regression_tests( # noqa: D417
test_timeout: int,
trace_id: str,
test_index: int,
is_numerical_code: bool | None = None, # noqa: FBT001
is_numerical_code: bool | None = None,
) -> tuple[str, str, str] | None:
"""Generate regression tests for the given function by making a request to the Django endpoint.

Expand Down Expand Up @@ -706,7 +706,7 @@ def generate_regression_tests( # noqa: D417
error = response.json()["error"]
logger.error(f"Error generating tests: {response.status_code} - {error}")
ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": error})
return None # noqa: TRY300
return None
except Exception:
logger.error(f"Error generating tests: {response.status_code} - {response.text}")
ph("cli-testgen-error-response", {"response_status_code": response.status_code, "error": response.text})
Expand All @@ -722,7 +722,7 @@ def get_optimization_review(
function_trace_id: str,
coverage_message: str,
replay_tests: str,
concolic_tests: str, # noqa: ARG002
concolic_tests: str,
calling_fn_details: str,
) -> OptimizationReviewResult:
"""Compute the optimization review of current Pull Request.
Expand Down
6 changes: 3 additions & 3 deletions codeflash/api/cfapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ def make_cfapi_request(
else:
response = requests.get(url, headers=cfapi_headers, params=params, timeout=60)
response.raise_for_status()
return response # noqa: TRY300
return response
except requests.exceptions.HTTPError:
# response may be either a string or JSON, so we handle both cases
error_message = ""
Expand All @@ -102,7 +102,7 @@ def make_cfapi_request(


@lru_cache(maxsize=1)
def get_user_id(api_key: Optional[str] = None) -> Optional[str]: # noqa: PLR0911
def get_user_id(api_key: Optional[str] = None) -> Optional[str]:
"""Retrieve the user's userid by making a request to the /cfapi/cli-get-user endpoint.

:param api_key: The API key to use. If None, uses get_codeflash_api_key().
Expand Down Expand Up @@ -396,7 +396,7 @@ def get_blocklisted_functions() -> dict[str, set[str]] | dict[str, Any]:

def is_function_being_optimized_again(
owner: str, repo: str, pr_number: int, code_contexts: list[dict[str, str]]
) -> Any: # noqa: ANN401
) -> Any:
"""Check if the function being optimized is being optimized again."""
response = make_cfapi_request(
"/is-already-optimized",
Expand Down
2 changes: 1 addition & 1 deletion codeflash/benchmarking/codeflash_trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ def __call__(self, func: Callable) -> Callable:
func_id = (func.__module__, func.__name__)

@functools.wraps(func)
def wrapper(*args, **kwargs) -> Any: # noqa: ANN002, ANN003, ANN401
def wrapper(*args, **kwargs) -> Any: # noqa: ANN002, ANN003
# Initialize thread-local active functions set if it doesn't exist
if not hasattr(self._thread_local, "active_functions"):
self._thread_local.active_functions = set()
Expand Down
2 changes: 1 addition & 1 deletion codeflash/benchmarking/instrument_codeflash_trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDe

return updated_node

def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: # noqa: ARG002
def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
# Create import statement for codeflash_trace
if not self.added_codeflash_trace:
return updated_node
Expand Down
12 changes: 6 additions & 6 deletions codeflash/benchmarking/plugin/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[Benchmark

# Process each row
for row in cursor.fetchall():
module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row
module_name, class_name, function_name, benchmark_file, benchmark_func, _benchmark_line, time_ns = row

# Create the function key (module_name.class_name.function_name)
if class_name:
Expand Down Expand Up @@ -172,7 +172,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:

# Process overhead information
for row in cursor.fetchall():
benchmark_file, benchmark_func, benchmark_line, total_overhead_ns = row
benchmark_file, benchmark_func, _benchmark_line, total_overhead_ns = row
benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
overhead_by_benchmark[benchmark_key] = total_overhead_ns or 0 # Handle NULL sum case

Expand All @@ -184,7 +184,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:

# Process each row and subtract overhead
for row in cursor.fetchall():
benchmark_file, benchmark_func, benchmark_line, time_ns = row
benchmark_file, benchmark_func, _benchmark_line, time_ns = row

# Create the benchmark key (file::function::line)
benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
Expand All @@ -200,7 +200,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:

# Pytest hooks
@pytest.hookimpl
def pytest_sessionfinish(self, session, exitstatus) -> None: # noqa: ANN001, ARG002
def pytest_sessionfinish(self, session, exitstatus) -> None:
"""Execute after whole test run is completed."""
# Write any remaining benchmark timings to the database
codeflash_trace.close()
Expand Down Expand Up @@ -236,7 +236,7 @@ class Benchmark: # noqa: D106
def __init__(self, request: pytest.FixtureRequest) -> None:
self.request = request

def __call__(self, func, *args, **kwargs): # type: ignore # noqa: ANN001, ANN002, ANN003, ANN204, PGH003
def __call__(self, func, *args, **kwargs): # type: ignore # noqa: ANN002, ANN003, ANN204, PGH003
"""Handle both direct function calls and decorator usage."""
if args or kwargs:
# Used as benchmark(func, *args, **kwargs)
Expand All @@ -249,7 +249,7 @@ def wrapped_func(*args, **kwargs): # noqa: ANN002, ANN003, ANN202
self._run_benchmark(func)
return wrapped_func

def _run_benchmark(self, func, *args, **kwargs): # noqa: ANN001, ANN002, ANN003, ANN202
def _run_benchmark(self, func, *args, **kwargs): # noqa: ANN002, ANN003, ANN202
"""Actual benchmark implementation."""
node_path = getattr(self.request.node, "path", None) or getattr(self.request.node, "fspath", None)
if node_path is None:
Expand Down
2 changes: 1 addition & 1 deletion codeflash/cli_cmds/cli_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def inquirer_wrapper(func: Callable[..., str | bool], *args: str | bool, **kwarg
return func(*new_args, **new_kwargs)


def split_string_to_cli_width(string: str, is_confirm: bool = False) -> list[str]: # noqa: FBT001, FBT002
def split_string_to_cli_width(string: str, is_confirm: bool = False) -> list[str]:
cli_width, _ = shutil.get_terminal_size()
# split string to lines that accommodate "[?] " prefix
cli_width -= len("[?] ")
Expand Down
19 changes: 7 additions & 12 deletions codeflash/cli_cmds/cmd_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -663,7 +663,7 @@ def create_empty_pyproject_toml(pyproject_toml_path: Path) -> None:
apologize_and_exit()


def install_github_actions(override_formatter_check: bool = False) -> None: # noqa: FBT001, FBT002
def install_github_actions(override_formatter_check: bool = False) -> None:
try:
config, _config_file_path = parse_config_file(override_formatter_check=override_formatter_check)

Expand Down Expand Up @@ -1089,11 +1089,12 @@ def install_github_actions(override_formatter_check: bool = False) -> None: # n
apologize_and_exit()


def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager: # noqa: PLR0911
def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager:
"""Determine which dependency manager is being used based on pyproject.toml contents."""
if (Path.cwd() / "poetry.lock").exists():
cwd = Path.cwd()
if (cwd / "poetry.lock").exists():
return DependencyManager.POETRY
if (Path.cwd() / "uv.lock").exists():
if (cwd / "uv.lock").exists():
Comment on lines +1094 to +1097
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 101% (1.01x) speedup for determine_dependency_manager in codeflash/cli_cmds/cmd_init.py

⏱️ Runtime : 3.86 milliseconds 1.91 milliseconds (best of 247 runs)

📝 Explanation and details

This optimization achieves a 101% speedup (2.02x faster) by replacing Python's pathlib.Path operations with lower-level os module functions for file system checks.

Key Changes

Replaced pathlib.Path with os module:

  • Path.cwd()os.getcwd() (12x faster: 17.0μs → 1.4μs per call)
  • (cwd / "poetry.lock").exists()os.path.exists(os.path.join(cwd, "poetry.lock")) (2.7x faster: 22.6μs → 8.3μs per call)
  • (cwd / "uv.lock").exists()os.path.exists(os.path.join(cwd, "uv.lock")) (2.9x faster: 20.4μs → 7.1μs per call)

Why This Is Faster

The line profiler shows 97% of execution time (20.2ms out of 20.8ms) was spent on just three lines:

  1. Getting current working directory (27.7% of total time)
  2. Checking poetry.lock existence (36.6% of total time)
  3. Checking uv.lock existence (32.7% of total time)

pathlib.Path provides a high-level, object-oriented interface with operator overloading (/) and method chaining, which adds overhead:

  • Each Path object creation involves class instantiation
  • The / operator triggers __truediv__ magic method calls
  • .exists() method has additional validation and type checking

In contrast, os module functions are thin wrappers around C-level system calls, avoiding Python object overhead entirely.

Impact on Real Workloads

Based on function_references, this function is called during GitHub Actions workflow generation in two hot paths:

  1. generate_dynamic_workflow_content() - Called when initializing CI/CD workflows
  2. customize_codeflash_yaml_content() - Called for workflow customization

Both functions execute dependency manager detection during repository setup. While setup is not a high-frequency operation, the 2x speedup improves developer experience during codeflash init runs.

Test Results Analysis

All test cases show 80-110% speedup, with particularly strong gains in:

  • Lock file detection tests: 85-109% faster (tests that hit the early-return path benefit most)
  • Large-scale tests (500 keys): 91% faster (showing the optimization scales well)
  • Mixed configuration tests: 86-104% faster (common in real projects)

The optimization is most effective when lock files exist (enabling early returns), but still provides significant gains even when the full tool section must be analyzed.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 240 Passed
⏪ Replay Tests 🔘 None Found
🔎 Concolic Coverage Tests 🔘 None Found
📊 Tests Coverage 100.0%
🌀 Click to see Generated Regression Tests
from enum import Enum

# imports
from codeflash.cli_cmds.cmd_init import determine_dependency_manager


# Provide the DependencyManager enum so the function under test can refer to it.
# In the real codebase this would live in the module being tested; for the purposes
# of this standalone test file we define it here with the exact members expected.
class DependencyManager(Enum):
    POETRY = "poetry"
    UV = "uv"
    PIP = "pip"
    UNKNOWN = "unknown"


def test_no_tool_key_defaults_to_pip(tmp_path, monkeypatch):
    # When pyproject.toml contains no "tool" section, the default is PIP.
    monkeypatch.chdir(tmp_path)  # ensure an empty cwd without lock files
    codeflash_output = determine_dependency_manager({})
    result = codeflash_output  # 36.5μs -> 19.7μs (85.2% faster)


def test_poetry_lock_file_takes_precedence_over_pyproject(tmp_path, monkeypatch):
    # If poetry.lock exists in cwd it should return POETRY regardless of pyproject contents.
    monkeypatch.chdir(tmp_path)
    # create poetry.lock file
    (tmp_path / "poetry.lock").write_text("")  # create an empty lock file
    # Even if pyproject indicates pip, the lock file should win
    pyproject = {"tool": {"pip": {}}}
    codeflash_output = determine_dependency_manager(pyproject)  # 21.9μs -> 10.4μs (109% faster)


def test_uv_lock_file_detected_when_no_poetry_lock(tmp_path, monkeypatch):
    # If uv.lock exists and poetry.lock does not, UV should be returned.
    monkeypatch.chdir(tmp_path)
    (tmp_path / "uv.lock").write_text("")  # create uv.lock
    codeflash_output = determine_dependency_manager({"tool": {}})  # 35.4μs -> 19.7μs (80.2% faster)


def test_poetry_key_in_tool_section_detected():
    # If "tool" contains a "poetry" key, detect POETRY.
    pyproject = {"tool": {"poetry": {"name": "example"}}}
    codeflash_output = determine_dependency_manager(pyproject)  # 33.9μs -> 17.8μs (90.9% faster)


def test_uv_key_in_tool_section_detected_before_pip():
    # Ensure UV is detected before PIP when both markers exist in tool section.
    pyproject = {"tool": {"pip": {}, "uv": {}}}
    # The function checks for poetry, then uv, then pip. Hence UV should be returned.
    codeflash_output = determine_dependency_manager(pyproject)  # 32.6μs -> 17.5μs (86.2% faster)


def test_uv_prefix_key_detected():
    # Keys that start with "uv" (e.g., "uv-core", "uvtool") should be detected as UV.
    pyproject = {"tool": {"uv-core": {}, "something": {}}}
    codeflash_output = determine_dependency_manager(pyproject)  # 32.3μs -> 17.3μs (86.7% faster)


def test_pip_detected_via_setuptools_key_too():
    # The presence of "setuptools" in tool should also map to PIP.
    pyproject = {"tool": {"setuptools": {"setup.cfg": {}}}}
    codeflash_output = determine_dependency_manager(pyproject)  # 32.0μs -> 17.1μs (87.3% faster)


def test_unknown_when_tool_present_but_no_markers():
    # If "tool" exists but doesn't contain poetry, uv*, pip, or setuptools, return UNKNOWN.
    pyproject = {"tool": {}}  # explicit empty tool section
    codeflash_output = determine_dependency_manager(pyproject)  # 32.0μs -> 16.9μs (89.4% faster)


def test_tool_value_as_string_poetry_behaviour():
    # If tool is a non-dict but contains the substring "poetry", the 'in' check will still
    # find it (since strings support membership checks). This is an edge case that follows
    # from the original implementation's use of 'in'.
    pyproject = {"tool": "poetry"}  # not typical but valid input shape to test behavior
    codeflash_output = determine_dependency_manager(pyproject)  # 31.7μs -> 17.0μs (86.6% faster)


def test_tool_value_as_string_uv_edge_case():
    # Another unusual shape: tool is a string that does not cause the uv detection to work
    # because the function iterates characters. This ensures we capture that behavior.
    pyproject = {"tool": "uvtool"}  # membership for poetry fails; any(...) will iterate chars
    # Expect UNKNOWN because 'poetry' not in "uvtool", the any(...) over characters won't startwith('uv'),
    # and neither 'pip' nor 'setuptools' are substrings that the function checks using 'in'.
    codeflash_output = determine_dependency_manager(pyproject)  # 31.8μs -> 17.1μs (85.7% faster)


def test_large_scale_many_keys_with_uv_at_end():
    # Large-scale style test: create many keys (but <1000) to ensure algorithm still finds uv
    # and does not time out or fail. We keep it deterministic and modest in size.
    many_keys = {f"key_{i}": {} for i in range(500)}  # 500 keys - within the requested limit
    many_keys["uv_final"] = {"config": True}  # add the uv-prefixed key at the end
    pyproject = {"tool": many_keys}
    codeflash_output = determine_dependency_manager(pyproject)  # 34.1μs -> 17.8μs (91.2% faster)


def test_ordering_poetry_over_uv_and_uv_over_pip(tmp_path, monkeypatch):
    # Create both lock files to ensure poetry.lock precedence over uv.lock
    monkeypatch.chdir(tmp_path)
    (tmp_path / "uv.lock").write_text("uv content")
    (tmp_path / "poetry.lock").write_text("poetry content")
    # poetry.lock checked first so expect POETRY
    codeflash_output = determine_dependency_manager({"tool": {"uv": {}, "pip": {}}})  # 22.5μs -> 11.0μs (104% faster)


def test_uv_detection_with_nonstandard_keys():
    # Ensure keys like 'uv123', 'uv-' still count as starting with 'uv'
    pyproject = {"tool": {"uv123": {}, "notuv": {}, "pip": {}}}
    codeflash_output = determine_dependency_manager(pyproject)  # 33.6μs -> 18.0μs (87.4% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
import os
import tempfile
from pathlib import Path

import pytest

from codeflash.cli_cmds.cmd_init import determine_dependency_manager


def test_poetry_lock_file_exists():
    """Test that poetry.lock file presence returns DependencyManager.POETRY."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create a poetry.lock file in the temporary directory
        poetry_lock_path = Path(tmpdir) / "poetry.lock"
        poetry_lock_path.touch()

        # Change to the temporary directory
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            codeflash_output = determine_dependency_manager({})
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_uv_lock_file_exists():
    """Test that uv.lock file presence returns DependencyManager.UV."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create a uv.lock file in the temporary directory
        uv_lock_path = Path(tmpdir) / "uv.lock"
        uv_lock_path.touch()

        # Change to the temporary directory
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            codeflash_output = determine_dependency_manager({})
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_poetry_in_tool_section():
    """Test that poetry in tool section returns DependencyManager.POETRY."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"poetry": {"name": "my-project", "version": "0.1.0"}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_uv_in_tool_section():
    """Test that uv tool config returns DependencyManager.UV."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"uv": {"project": {"requires-python": ">=3.8"}}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_pip_in_tool_section():
    """Test that pip in tool section returns DependencyManager.PIP."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"pip": {"index-url": "https://pypi.org/simple"}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_setuptools_in_tool_section():
    """Test that setuptools in tool section returns DependencyManager.PIP."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"setuptools": {"packages": ["my_package"]}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_no_tool_section_returns_pip():
    """Test that missing tool section defaults to DependencyManager.PIP."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"build-system": {"requires": ["setuptools"]}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_empty_pyproject_returns_pip():
    """Test that empty pyproject.toml data defaults to DependencyManager.PIP."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            codeflash_output = determine_dependency_manager({})
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_poetry_lock_takes_precedence_over_pyproject():
    """Test that poetry.lock file takes precedence over pyproject.toml tool section."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create poetry.lock file
        poetry_lock_path = Path(tmpdir) / "poetry.lock"
        poetry_lock_path.touch()

        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Even with uv in tool section, poetry.lock should take precedence
            pyproject_data = {"tool": {"uv": {}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_uv_lock_takes_precedence_over_pyproject():
    """Test that uv.lock file takes precedence over pyproject.toml tool section."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create uv.lock file
        uv_lock_path = Path(tmpdir) / "uv.lock"
        uv_lock_path.touch()

        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Even with poetry in tool section, uv.lock should take precedence
            pyproject_data = {"tool": {"poetry": {}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_poetry_lock_takes_precedence_over_uv_lock():
    """Test that poetry.lock is checked before uv.lock."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create both lock files
        poetry_lock_path = Path(tmpdir) / "poetry.lock"
        poetry_lock_path.touch()
        uv_lock_path = Path(tmpdir) / "uv.lock"
        uv_lock_path.touch()

        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            codeflash_output = determine_dependency_manager({})
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_tool_section_exists_but_empty():
    """Test that empty tool section returns DependencyManager.UNKNOWN."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_unknown_tool_in_tool_section():
    """Test that unknown tool in tool section returns DependencyManager.UNKNOWN."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"black": {"line-length": 88}, "isort": {"profile": "black"}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_uv_tool_with_prefix_matching():
    """Test that tools starting with 'uv' are recognized."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Test with 'uv-config' which starts with 'uv'
            pyproject_data = {"tool": {"uv-config": {"some-setting": "value"}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_mixed_tools_poetry_takes_precedence():
    """Test that poetry takes precedence when both poetry and other tools are present."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"poetry": {"name": "project"}, "black": {"line-length": 88}, "setuptools": {}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_mixed_tools_uv_over_pip():
    """Test that uv takes precedence over pip when both are present."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"uv": {}, "pip": {}, "setuptools": {}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_tool_section_with_none_value():
    """Test handling of tool section with None value."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": None}
            # This should raise an error or be handled gracefully
            # The function iterates over tool_section keys, which would fail with None
            with pytest.raises((TypeError, AttributeError)):
                determine_dependency_manager(pyproject_data)
        finally:
            os.chdir(original_cwd)


def test_case_sensitive_tool_names():
    """Test that tool names are case-sensitive."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # 'Poetry' with capital P should not match 'poetry'
            pyproject_data = {"tool": {"Poetry": {"name": "project"}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_uv_as_substring_check():
    """Test that 'uv' prefix check works correctly."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Tool name starting with 'uv' should be recognized
            pyproject_data = {"tool": {"uvicorn": {}}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_nonexistent_directory_handling():
    """Test behavior when checking for lock files in nonexistent directory."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            # Create a path and immediately delete it
            test_path = Path(tmpdir) / "subdir"
            test_path.mkdir()
            os.chdir(test_path)
            test_path.rmdir()

            # This will try to check for lock files in a deleted directory
            # The Path.exists() should handle this gracefully
            with pytest.raises(FileNotFoundError):
                os.chdir(test_path)
        finally:
            os.chdir(original_cwd)


def test_large_tool_section_with_many_tools():
    """Test performance with a large number of tools in tool section."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Create a tool section with many tools
            tool_data = {f"tool_{i}": {"setting": "value"} for i in range(100)}
            tool_data["poetry"] = {"name": "project"}

            pyproject_data = {"tool": tool_data}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_large_pyproject_with_many_sections():
    """Test performance with a large pyproject.toml data structure."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Create a large pyproject.toml data structure
            pyproject_data = {
                "build-system": {"requires": ["poetry-core"]},
                "project": {"name": "my-project", "dependencies": [f"dep_{i}" for i in range(100)]},
                "tool": {f"tool_{i}": {"config": "value"} for i in range(50)},
            }
            pyproject_data["tool"]["uv"] = {}

            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_deep_nested_pyproject_structure():
    """Test with deeply nested pyproject.toml structure."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Create a deeply nested structure
            pyproject_data = {
                "tool": {"poetry": {"group": {"dev": {"dependencies": {"pytest": "^7.0", "black": "^22.0"}}}}}
            }
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_multiple_uv_prefixed_tools():
    """Test recognition of multiple tools starting with 'uv' prefix."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Multiple tools with 'uv' prefix
            pyproject_data = {"tool": {f"uv_tool_{i}": {"config": "value"} for i in range(50)}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_large_tool_data_with_many_keys():
    """Test with large amount of data in tool configurations."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Create a large tool configuration
            setuptools_config = {f"package_{i}": f"config_{i}" for i in range(200)}
            pyproject_data = {"tool": {"setuptools": setuptools_config}}
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_mixed_large_structure_poetry_takes_precedence():
    """Test with large mixed structure where poetry should win."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Create a large structure with many tools
            pyproject_data = {
                "project": {"dependencies": [f"dep_{i}" for i in range(50)]},
                "tool": {
                    "black": {f"setting_{i}": "value" for i in range(20)},
                    "isort": {f"setting_{i}": "value" for i in range(20)},
                    "pytest": {f"setting_{i}": "value" for i in range(20)},
                    "poetry": {"name": "project", "version": "0.1.0"},
                    "setuptools": {f"setting_{i}": "value" for i in range(20)},
                },
            }
            codeflash_output = determine_dependency_manager(pyproject_data)
            result = codeflash_output
        finally:
            os.chdir(original_cwd)


def test_performance_with_many_iterations():
    """Test that the function performs consistently across multiple calls."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            pyproject_data = {"tool": {"uv": {}}}

            # Call the function multiple times and verify consistent results
            results = [determine_dependency_manager(pyproject_data) for _ in range(100)]
        finally:
            os.chdir(original_cwd)


def test_stress_test_many_lock_files_check():
    """Test performance when checking for lock files that don't exist."""
    with tempfile.TemporaryDirectory() as tmpdir:
        original_cwd = Path.cwd()
        try:
            os.chdir(tmpdir)
            # Call the function many times with no lock files
            for _ in range(200):
                codeflash_output = determine_dependency_manager({})
                result = codeflash_output
        finally:
            os.chdir(original_cwd)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To test or edit this optimization locally git merge codeflash/optimize-pr1195-2026-01-29T17.28.26

Suggested change
cwd = Path.cwd()
if (cwd / "poetry.lock").exists():
return DependencyManager.POETRY
if (Path.cwd() / "uv.lock").exists():
if (cwd / "uv.lock").exists():
cwd = os.getcwd()
if os.path.exists(os.path.join(cwd, "poetry.lock")):
return DependencyManager.POETRY
if os.path.exists(os.path.join(cwd, "uv.lock")):

return DependencyManager.UV
if "tool" not in pyproject_data:
return DependencyManager.PIP
Expand Down Expand Up @@ -1251,10 +1252,7 @@ def collect_repo_files_for_workflow(git_root: Path) -> dict[str, Any]:


def generate_dynamic_workflow_content(
optimize_yml_content: str,
config: tuple[dict[str, Any], Path],
git_root: Path,
benchmark_mode: bool = False, # noqa: FBT001, FBT002
optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False
) -> str:
"""Generate workflow content with dynamic steps from AI service, falling back to static template.

Expand Down Expand Up @@ -1378,10 +1376,7 @@ def generate_dynamic_workflow_content(


def customize_codeflash_yaml_content(
optimize_yml_content: str,
config: tuple[dict[str, Any], Path],
git_root: Path,
benchmark_mode: bool = False, # noqa: FBT001, FBT002
optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False
) -> str:
module_path = str(Path(config["module_root"]).relative_to(git_root) / "**")
optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path)
Expand Down
54 changes: 26 additions & 28 deletions codeflash/code_utils/code_extractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,28 +141,28 @@ def __init__(self) -> None:
self.scope_depth = 0
self.if_else_depth = 0

def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]: # noqa: ARG002
def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:
self.scope_depth += 1
return True

def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None: # noqa: ARG002
def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:
self.scope_depth -= 1

def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]: # noqa: ARG002
def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:
self.scope_depth += 1
return True

def leave_ClassDef(self, original_node: cst.ClassDef) -> None: # noqa: ARG002
def leave_ClassDef(self, original_node: cst.ClassDef) -> None:
self.scope_depth -= 1

def visit_If(self, node: cst.If) -> Optional[bool]: # noqa: ARG002
def visit_If(self, node: cst.If) -> Optional[bool]:
self.if_else_depth += 1
return True

def leave_If(self, original_node: cst.If) -> None: # noqa: ARG002
def leave_If(self, original_node: cst.If) -> None:
self.if_else_depth -= 1

def visit_Else(self, node: cst.Else) -> Optional[bool]: # noqa: ARG002
def visit_Else(self, node: cst.Else) -> Optional[bool]:
# Else blocks are already counted as part of the if statement
return True

Expand Down Expand Up @@ -231,24 +231,24 @@ def __init__(self, new_assignments: dict[str, cst.Assign | cst.AnnAssign], new_a
self.scope_depth = 0
self.if_else_depth = 0

def visit_FunctionDef(self, node: cst.FunctionDef) -> None: # noqa: ARG002
def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
self.scope_depth += 1

def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: # noqa: ARG002
def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
self.scope_depth -= 1
return updated_node

def visit_ClassDef(self, node: cst.ClassDef) -> None: # noqa: ARG002
def visit_ClassDef(self, node: cst.ClassDef) -> None:
self.scope_depth += 1

def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: # noqa: ARG002
def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:
self.scope_depth -= 1
return updated_node

def visit_If(self, node: cst.If) -> None: # noqa: ARG002
def visit_If(self, node: cst.If) -> None:
self.if_else_depth += 1

def leave_If(self, original_node: cst.If, updated_node: cst.If) -> cst.If: # noqa: ARG002
def leave_If(self, original_node: cst.If, updated_node: cst.If) -> cst.If:
self.if_else_depth -= 1
return updated_node

Expand Down Expand Up @@ -283,7 +283,7 @@ def leave_AnnAssign(self, original_node: cst.AnnAssign, updated_node: cst.AnnAss

return updated_node

def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: # noqa: ARG002
def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
# Add any new assignments that weren't in the original file
new_statements = list(updated_node.body)

Expand Down Expand Up @@ -396,20 +396,20 @@ def __init__(self) -> None:
self.global_statements = []
self.in_function_or_class = False

def visit_ClassDef(self, node: cst.ClassDef) -> bool: # noqa: ARG002
def visit_ClassDef(self, node: cst.ClassDef) -> bool:
# Don't visit inside classes
self.in_function_or_class = True
return False

def leave_ClassDef(self, original_node: cst.ClassDef) -> None: # noqa: ARG002
def leave_ClassDef(self, original_node: cst.ClassDef) -> None:
self.in_function_or_class = False

def visit_FunctionDef(self, node: cst.FunctionDef) -> bool: # noqa: ARG002
def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
# Don't visit inside functions
self.in_function_or_class = True
return False

def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None: # noqa: ARG002
def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:
self.in_function_or_class = False

def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
Expand Down Expand Up @@ -490,16 +490,16 @@ def visit_Module(self, node: cst.Module) -> None:
self.depth = 0
self._collect_imports_from_block(node)

def visit_FunctionDef(self, node: cst.FunctionDef) -> None: # noqa: ARG002
def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
self.depth += 1

def leave_FunctionDef(self, node: cst.FunctionDef) -> None: # noqa: ARG002
def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
self.depth -= 1

def visit_ClassDef(self, node: cst.ClassDef) -> None: # noqa: ARG002
def visit_ClassDef(self, node: cst.ClassDef) -> None:
self.depth += 1

def leave_ClassDef(self, node: cst.ClassDef) -> None: # noqa: ARG002
def leave_ClassDef(self, node: cst.ClassDef) -> None:
self.depth -= 1

def visit_If(self, node: cst.If) -> None:
Expand Down Expand Up @@ -529,9 +529,7 @@ def find_last_import_line(target_code: str) -> int:

class FutureAliasedImportTransformer(cst.CSTTransformer):
def leave_ImportFrom(
self,
original_node: cst.ImportFrom, # noqa: ARG002
updated_node: cst.ImportFrom,
self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
) -> cst.BaseSmallStatement | cst.FlattenSentinel[cst.BaseSmallStatement] | cst.RemovalSentinel:
import libcst.matchers as m

Expand Down Expand Up @@ -676,7 +674,7 @@ def resolve_star_import(module_name: str, project_root: Path) -> set[str]:
if not name.startswith("_"):
public_names.add(name)

return public_names # noqa: TRY300
return public_names

except Exception as e:
logger.warning(f"Error resolving star import for {module_name}: {e}")
Expand Down Expand Up @@ -1165,7 +1163,7 @@ def _is_target_function_call(self, node: ast.Call) -> bool:

return False

def _get_call_name(self, func_node) -> Optional[str]: # noqa: ANN001
def _get_call_name(self, func_node) -> Optional[str]:
"""Extract the name being called from a function node."""
# Fast path short-circuit for ast.Name nodes
if isinstance(func_node, ast.Name):
Expand Down Expand Up @@ -1555,7 +1553,7 @@ def is_numerical_code(code_string: str, function_name: str | None = None) -> boo

# If numba is not installed and all modules used require numba for optimization,
# return False since we can't optimize this code
if not has_numba and modules_used.issubset(NUMBA_REQUIRED_MODULES): # noqa : SIM103
if not has_numba and modules_used.issubset(NUMBA_REQUIRED_MODULES):
return False

return True
Expand Down
Loading
Loading