diff --git a/codeflash/languages/code_replacer.py b/codeflash/languages/code_replacer.py index f9853fbfc..56d5fd64f 100644 --- a/codeflash/languages/code_replacer.py +++ b/codeflash/languages/code_replacer.py @@ -6,13 +6,16 @@ from __future__ import annotations -from pathlib import Path +import logging +import os from typing import TYPE_CHECKING from codeflash.cli_cmds.console import logger from codeflash.languages.base import FunctionFilterCriteria, Language if TYPE_CHECKING: + from pathlib import Path + from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.base import LanguageSupport from codeflash.models.models import CodeStringsMarkdown @@ -27,7 +30,8 @@ def get_optimized_code_for_module( from codeflash.languages.current import is_python file_to_code_context = optimized_code.file_to_path() - module_optimized_code = file_to_code_context.get(str(relative_path)) + target_key = str(relative_path) + module_optimized_code = file_to_code_context.get(target_key) if module_optimized_code is not None: return module_optimized_code @@ -36,23 +40,36 @@ def get_optimized_code_for_module( # Fallback 1: single code block with no file path if "None" in file_to_code_context and len(file_to_code_context) == 1: - logger.debug(f"Using code block with None file_path for {relative_path}") + if logger.isEnabledFor(logging.DEBUG): + logger.debug(f"Using code block with None file_path for {relative_path}") return file_to_code_context["None"] # Fallback 2: match by filename (basename) — the LLM sometimes returns a different # directory prefix but the correct filename target_name = relative_path.name - basename_matches = [ - code for path, code in file_to_code_context.items() if path != "None" and Path(path).name == target_name - ] - if len(basename_matches) == 1: - logger.debug(f"Using basename-matched code block for {relative_path}") - return basename_matches[0] + # Avoid building a full list of matches and avoid Path() construction for each key. + first_match: str = "" + match_count = 0 + for path, code in file_to_code_context.items(): + if path == "None": + continue + # os.path.basename is faster than constructing Path objects for each key + if os.path.basename(path) == target_name: # noqa: PTH119 + first_match = code + match_count += 1 + if match_count > 1: + break + + if match_count == 1: + if logger.isEnabledFor(logging.DEBUG): + logger.debug(f"Using basename-matched code block for {relative_path}") + return first_match # Fallback 3: single code block for non-Python (AI often returns one block with wrong path) if len(file_to_code_context) == 1 and not is_python(): only_key = next(iter(file_to_code_context.keys())) - logger.debug(f"Using only code block {only_key} for {relative_path}") + if logger.isEnabledFor(logging.DEBUG): + logger.debug(f"Using only code block {only_key} for {relative_path}") return file_to_code_context[only_key] logger.warning(