⚡️ Speed up function get_optimized_code_for_module by 12% in PR #1660 (unstructured-inference)

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}"

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: