diff --git a/codeflash/languages/python/static_analysis/code_replacer.py b/codeflash/languages/python/static_analysis/code_replacer.py index c230dfacd..26b6f1e6d 100644 --- a/codeflash/languages/python/static_analysis/code_replacer.py +++ b/codeflash/languages/python/static_analysis/code_replacer.py @@ -1,6 +1,7 @@ from __future__ import annotations import ast +import os from collections import defaultdict from functools import lru_cache from itertools import chain @@ -557,25 +558,31 @@ def _extract_function_from_code( def get_optimized_code_for_module(relative_path: Path, optimized_code: CodeStringsMarkdown) -> str: file_to_code_context = optimized_code.file_to_path() - module_optimized_code = file_to_code_context.get(str(relative_path)) + str_relative = str(relative_path) + module_optimized_code = file_to_code_context.get(str_relative) if module_optimized_code is None: # Fallback: if there's only one code block with None file path, # use it regardless of the expected path (the AI server doesn't always include file paths) - if "None" in file_to_code_context and len(file_to_code_context) == 1: + if len(file_to_code_context) == 1 and "None" in file_to_code_context: module_optimized_code = file_to_code_context["None"] logger.debug(f"Using code block with None file_path for {relative_path}") else: - available_files = list(file_to_code_context.keys()) # Check if this looks like a path mismatch (same filename exists under a different path) # vs the AI simply not returning code for this module requested_name = relative_path.name - similar = [f for f in available_files if Path(f).name == requested_name] + similar: list[str] = [] + # Iterate keys once and avoid constructing Path objects repeatedly + for f in file_to_code_context: + if os.path.basename(f) == requested_name: + similar.append(f) if similar: + available_files = list(file_to_code_context.keys()) logger.warning( f"Optimized code not found for '{relative_path}' but found similar path(s): {similar}. " f"Re-check your markdown code structure. Available files: {available_files}" ) else: + available_files = list(file_to_code_context.keys()) logger.debug( f"AI service did not return optimized code for '{relative_path}'. " f"Available files in response: {available_files}" diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 697601403..e955c94f7 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -331,12 +331,13 @@ def file_to_path(self) -> dict[str, str]: dict[str, str]: Mapping from file path (as string) to code. """ - if self._cache.get("file_to_path") is not None: + try: return self._cache["file_to_path"] - self._cache["file_to_path"] = { - str(code_string.file_path): code_string.code for code_string in self.code_strings - } - return self._cache["file_to_path"] + except KeyError: + # Build the mapping once and cache it + mapping = {str(code_string.file_path): code_string.code for code_string in self.code_strings} + self._cache["file_to_path"] = mapping + return mapping @staticmethod def parse_markdown_code(markdown_code: str, expected_language: str = "python") -> CodeStringsMarkdown: