fix: lazy cross module base class imports for python speedups

packages/@jsii/python-runtime/src/jsii/__init__.py

@@ -14,6 +14,7 @@
     kernel,
     proxy_for,
 )
+from ._utils import _LazyImport, _memoized
 from . import python
 
 
@@ -49,6 +50,8 @@
 __all__ = [
     "__version__",
     "__jsii_runtime_version__",
+    "_LazyImport",
+    "_memoized",
     "JSIIAssembly",
     "JSIIMeta",
     "JSIIAbstractClass",

packages/@jsii/python-runtime/src/jsii/_reference_map.py

@@ -1,6 +1,7 @@
 # This module exists to break an import cycle between jsii.runtime and jsii.kernel
 import importlib
 import inspect
+import sys
 
 from typing import Any, Iterable, List, Mapping, MutableMapping, Type
 from ._kernel.types import ObjRef
@@ -38,33 +39,25 @@ def register_interface(iface: Any):
     _interfaces[iface.__jsii_type__] = iface
 
 
-def _try_import_type_module(class_fqn: str) -> bool:
-    """Attempt to import the Python module containing a jsii type by FQN.
+def _find_containing_module(class_fqn: str) -> "tuple[str, list[str]] | None":
+    """Find the Python module that should contain the given jsii type.
 
-    With PEP 562 lazy loading, submodules are not imported until first attribute
-    access. If the jsii runtime needs to deserialize a type from a submodule
-    that hasn't been imported yet (e.g., a callback returns an object whose type
-    lives in an unloaded submodule), this function triggers the import so that
-    the type self-registers with the runtime.
+    Returns a tuple of (python_module_name, remaining_type_parts) if found,
+    or None if no module could be identified.
 
-    Resolution uses the ``_submodule_fqn_map`` (populated at assembly load time
-    from the code-generated ``_SUBMODULE_FQN_MAP`` dict) to deterministically
-    find the correct Python module for any FQN, even when submodules have custom
-    Python target names that differ from the FQN structure.
-
-    Returns True if an import was successfully triggered, False otherwise.
+    Resolution strategy:
+    1. Use the submodule FQN map (most-specific to least-specific prefix)
+    2. Fall back to the assembly's root module
     """
-    # Split FQN into components: e.g. "jsii-calc.cdk16625.donotimport.MyType"
     parts = class_fqn.split(".")
     if len(parts) < 2:
-        return False
+        return None
 
-    # The first component is the assembly name (may contain special chars for
-    # scoped packages, but the FQN uses it as-is).
+    # The first component is the assembly name
     assembly_name = parts[0]
     root_module = _assembly_to_module.get(assembly_name)
     if root_module is None:
-        return False
+        return None
 
     # Strategy 1: Use the submodule FQN map for deterministic resolution.
     # Walk from most-specific to least-specific prefix of the FQN to find
@@ -77,21 +70,103 @@ def _try_import_type_module(class_fqn: str) -> bool:
             if python_module is not None:
                 try:
                     importlib.import_module(python_module)
-                    return True
+                    return (python_module, parts[depth:])
                 except (ImportError, ModuleNotFoundError):
                     continue
 
     # Strategy 2: Fall back to importing the root module itself.
     # The type might live at the assembly root (no submodule).
     try:
         importlib.import_module(root_module)
-        return True
+        return (root_module, parts[1:])
     except (ImportError, ModuleNotFoundError):
-        pass
+        return None
+
+
+def _resolve_type_on_module(python_module: str, type_parts: "list[str]") -> bool:
+    """Traverse a module using dot-separated type parts to trigger lazy loading.
+
+    For eagerly-loaded modules, the type self-registers during import via the
+    JSIIMeta metaclass. For lazily-loaded modules, accessing an attribute on
+    the module triggers ``__getattr__`` which calls the factory function.
+
+    For nested types like "Parent.Nested", accessing "Parent" on the module
+    triggers the factory which defines both Parent and Nested inside it.
+
+    Returns True if the traversal succeeded (type should now be registered).
+    """
+    mod = sys.modules.get(python_module)
+    if mod is None or not type_parts:
+        return False
+
+    try:
+        obj = mod
+        for part in type_parts:
+            obj = getattr(obj, part)
+        return True
+    except AttributeError:
+        return False
+
+
+def _try_import_type_module(class_fqn: str) -> bool:
+    """Attempt to import the Python module containing a jsii type by FQN.
+
+    With PEP 562 lazy loading, submodules are not imported until first attribute
+    access. If the jsii runtime needs to deserialize a type from a submodule
+    that hasn't been imported yet (e.g., a callback returns an object whose type
+    lives in an unloaded submodule), this function triggers the import so that
+    the type self-registers with the runtime.
+
+    Resolution uses the ``_submodule_fqn_map`` (populated at assembly load time
+    from the code-generated ``_SUBMODULE_FQN_MAP`` dict) to deterministically
+    find the correct Python module for any FQN, even when submodules have custom
+    Python target names that differ from the FQN structure.
+
+    After a successful import, if the type is still not registered (because the
+    module uses lazy class factories), this function triggers the factory by
+    calling ``getattr()`` progressively on the module with each component of
+    the type's name path.
+
+    Returns True if an import was successfully triggered, False otherwise.
+    """
+    # Phase 1: Find the containing module and import it.
+    result = _find_containing_module(class_fqn)
+    if result is None:
+        return False
+
+    python_module, type_parts = result
+
+    # If the type is already registered (eager code path), we're done.
+    if class_fqn in _types or class_fqn in _data_types or class_fqn in _enums:
+        return True
+
+    # Phase 2: Type not registered yet — traverse the module to trigger
+    # lazy factories, then check registration.
+    if _resolve_type_on_module(python_module, type_parts):
+        return class_fqn in _types or class_fqn in _data_types or class_fqn in _enums
 
     return False
 
 
+def _obtain_interface(fqn: str) -> Any:
+    """Look up an interface by FQN, triggering lazy loading if necessary.
+
+    Returns the interface class. Raises ValueError if the interface cannot
+    be found after attempting lazy resolution.
+    """
+    iface = _interfaces.get(fqn)
+    if iface is not None:
+        return iface
+
+    _try_import_type_module(fqn)
+
+    iface = _interfaces.get(fqn)
+    if iface is not None:
+        return iface
+
+    raise ValueError(f"Unknown interface: {fqn}")
+
+
 class _FakeReference:
     def __init__(self, ref: str) -> None:
         self.__jsii_ref__ = ref
@@ -167,6 +242,11 @@ def resolve(self, kernel, ref):
             return _enums[class_fqn]
         elif class_fqn == "Object":
             # If any one interface is a struct, all of them are guaranteed to be (Kernel invariant)
+            # Ensure all interfaces/structs are loaded (may be behind lazy factories)
+            if ref.interfaces is not None:
+                for fqn in ref.interfaces:
+                    if fqn not in _data_types and fqn not in _interfaces:
+                        _try_import_type_module(fqn)
             if ref.interfaces is not None and any(
                 fqn in _data_types for fqn in ref.interfaces
             ):
@@ -210,7 +290,7 @@ def resolve_id(self, id: str) -> Any:
         return self._refs[id]
 
     def build_interface_proxies_for_ref(self, ref: ObjRef) -> List[Any]:
-        ifaces = [_interfaces[fqn] for fqn in ref.interfaces or []]
+        ifaces = [_obtain_interface(fqn) for fqn in ref.interfaces or []]
         classes = [iface.__jsii_proxy_class__() for iface in ifaces]
 
         # If there's no classes, use an Opaque reference to make sure the

packages/@jsii/python-runtime/src/jsii/_utils.py

@@ -1,4 +1,5 @@
 import functools
+import importlib
 
 from typing import Any, MutableMapping, Type
 
@@ -24,3 +25,47 @@ def wrapped(self):
         return stored[0]
 
     return property(wrapped)
+
+
+def _memoized(func):
+    """Cache the return value of a no-arg factory function."""
+    sentinel = object()
+    result = sentinel
+
+    def wrapper():
+        nonlocal result
+        if result is sentinel:
+            result = func()
+        return result
+
+    wrapper.__name__ = func.__name__
+    wrapper.__wrapped__ = func  # type: ignore[attr-defined]
+    return wrapper
+
+
+class _LazyImport:
+    """Defers ``importlib.import_module()`` until first attribute access.
+
+    This is used by jsii-pacmak generated code to lazily import cross-module
+    dependencies, breaking circular import chains and reducing module load time.
+
+    The imported module is cached after first successful resolution. Failed
+    imports are NOT cached, allowing retry on subsequent access.
+    """
+
+    def __init__(self, module_name: str, package: str | None = None) -> None:
+        self._module_name = module_name
+        self._package = package
+        self._module: Any = None
+
+    def __getattr__(self, name: str) -> Any:
+        if self._module is None:
+            self._module = importlib.import_module(self._module_name, self._package)
+        return getattr(self._module, name)
+
+    def __repr__(self) -> str:
+        if self._module is not None:
+            return repr(self._module)
+        if self._package:
+            return f"_LazyImport({self._module_name!r}, {self._package!r})"
+        return f"_LazyImport({self._module_name!r})"