⚡️ Speed up method JavaSupport.format_code by 36% in PR #1199 (omni-java)

codeflash/languages/java/formatter.py

@@ -14,6 +14,8 @@
 import tempfile
 from pathlib import Path
 
+_FORMATTER_CACHE: dict[str | None, JavaFormatter] = {}
+
 logger = logging.getLogger(__name__)
 
 
@@ -236,7 +238,7 @@ def format_java_code(source: str, project_root: Path | None = None) -> str:
         Formatted source code.
 
     """
-    formatter = JavaFormatter(project_root)
+    formatter = _get_cached_formatter(project_root)
     return formatter.format_code(source)
 
 
@@ -327,3 +329,12 @@ def normalize_java_code(source: str) -> str:
             normalized_lines.append(stripped)
 
     return "\n".join(normalized_lines)
+
+
+def _get_cached_formatter(project_root: Path | None) -> JavaFormatter:
+    key = str(project_root) if project_root is not None else None
+    fmt = _FORMATTER_CACHE.get(key)
+    if fmt is None:
+        fmt = JavaFormatter(project_root)
+        _FORMATTER_CACHE[key] = fmt
+    return fmt

codeflash/languages/registry.py

@@ -50,13 +50,13 @@ def _ensure_languages_registered() -> None:
     import contextlib
 
     with contextlib.suppress(ImportError):
-        from codeflash.languages.python import support as _
+        from codeflash.languages.python import support as _python_support  # noqa: F401
 
     with contextlib.suppress(ImportError):
-        from codeflash.languages.javascript import support as _
+        from codeflash.languages.javascript import support as _js_support  # noqa: F401
 
     with contextlib.suppress(ImportError):
-        from codeflash.languages.java import support as _
+        from codeflash.languages.java import support as _java_support  # noqa: F401
 
     _languages_registered = True