⚡️ Speed up function __getattr__ by 34% in PR #1887 (codeflash_python)

codeflash/languages/__init__.py

@@ -17,6 +17,8 @@
     new_source = lang.replace_function(file_path, function, new_code)
 """
 
+import importlib
+
 from codeflash.languages.base import (
     CodeContext,
     DependencyResolver,
@@ -60,29 +62,36 @@
     set_current_test_framework,
 )
 
+_LAZY_IMPORTS: dict[str, tuple[str, str]] = {
+    "FunctionInfo": ("codeflash_core.models", "FunctionToOptimize"),
+    "JavaScriptSupport": ("codeflash.languages.javascript.support", "JavaScriptSupport"),
+    "TypeScriptSupport": ("codeflash.languages.javascript.support", "TypeScriptSupport"),
+    "PythonSupport": ("codeflash.languages.python.support", "PythonSupport"),
+    "JavaSupport": ("codeflash.languages.java.support", "JavaSupport"),
+}
 
-# Lazy imports to avoid circular imports
-def __getattr__(name: str):
-    if name == "FunctionInfo":
-        from codeflash_core.models import FunctionToOptimize
-
-        return FunctionToOptimize
-    if name == "JavaScriptSupport":
-        from codeflash.languages.javascript.support import JavaScriptSupport
+_LAZY_CACHE: dict[str, object] = {}
 
-        return JavaScriptSupport
-    if name == "TypeScriptSupport":
-        from codeflash.languages.javascript.support import TypeScriptSupport
 
-        return TypeScriptSupport
-    if name == "PythonSupport":
-        from codeflash.languages.python.support import PythonSupport
-
-        return PythonSupport
-    if name == "JavaSupport":
-        from codeflash.languages.java.support import JavaSupport
+# Lazy imports to avoid circular imports
+def __getattr__(name: str):
+    # Fast path: return cached value if already resolved
+    try:
+        return _LAZY_CACHE[name]
+    except KeyError:
+        pass
+
+    importer = _LAZY_IMPORTS.get(name)
+    if importer is not None:
+        module_path, attr_name = importer
+        module = importlib.import_module(module_path)
+        value = getattr(module, attr_name)
+        # Cache both in module globals (so future attribute access avoids __getattr__)
+        # and in the local cache to speed subsequent lookups in this function.
+        globals()[name] = value
+        _LAZY_CACHE[name] = value
+        return value
 
-        return JavaSupport
     msg = f"module {__name__!r} has no attribute {name!r}"
     raise AttributeError(msg)