feat: replace eager submodule imports with lazy loading

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

@@ -1,4 +1,5 @@
 # This module exists to break an import cycle between jsii.runtime and jsii.kernel
+import importlib
 import inspect
 
 from typing import Any, Iterable, List, Mapping, MutableMapping, Type
@@ -10,6 +11,11 @@
 _enums: MutableMapping[str, Any] = {}
 _interfaces: MutableMapping[str, Any] = {}
 
+# Mapping from jsii assembly name to Python root module name, populated by
+# JSIIAssembly.load() so that on-demand type resolution can trigger imports
+# of lazily-loaded submodules when the kernel returns an unknown type.
+_assembly_to_module: MutableMapping[str, str] = {}
+
 
 def register_type(klass: Type):
     _types[klass.__jsii_type__] = klass
@@ -27,6 +33,48 @@ 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.
+
+    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.
+
+    The FQN format is: ``assembly_name.submodule.path.TypeName``
+    We strip the type name (last dot-separated component) and try to import
+    progressively shorter module paths until one succeeds.
+
+    Returns True if an import was successfully triggered, False otherwise.
+    """
+    # Split FQN into components: e.g. "jsii-calc.cdk16625.donotimport.MyType"
+    parts = class_fqn.split(".")
+    if len(parts) < 2:
+        return False
+
+    # 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
+
+    # Try importing submodule paths from most specific to least specific
+    # e.g. for "jsii-calc.cdk16625.donotimport.MyType":
+    #   try: jsii_calc.cdk16625.donotimport
+    #   try: jsii_calc.cdk16625
+    submodule_parts = parts[1:]  # Remove assembly name
+    for depth in range(len(submodule_parts), 0, -1):
+        module_path = f"{root_module}.{'.'.join(submodule_parts[:depth])}"
+        try:
+            importlib.import_module(module_path)
+            return True
+        except (ImportError, ModuleNotFoundError):
+            continue
+
+    return False
+
+
 class _FakeReference:
     def __init__(self, ref: str) -> None:
         self.__jsii_ref__ = ref
@@ -128,6 +176,17 @@ def resolve(self, kernel, ref):
             else:
                 return InterfaceDynamicProxy(self.build_interface_proxies_for_ref(ref))
         else:
+            # The type isn't registered yet. With lazy loading, it may live in
+            # a submodule that hasn't been imported. Try importing the module
+            # that should contain this type — importing triggers type
+            # registration as a side effect — then retry the lookup.
+            if _try_import_type_module(class_fqn):
+                if class_fqn in _types:
+                    return self.resolve(kernel, ref)
+                elif class_fqn in _data_types:
+                    return self.resolve(kernel, ref)
+                elif class_fqn in _enums:
+                    return self.resolve(kernel, ref)
             raise ValueError(f"Unknown type: {class_fqn}")
 
     def resolve_id(self, id: str) -> Any:

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

@@ -54,6 +54,10 @@ def load(cls, *args, _kernel=kernel, **kwargs) -> "JSIIAssembly":
         ) as assembly_path:
             _kernel.load(assembly.name, assembly.version, os.fspath(assembly_path))
 
+        # Register the assembly-to-module mapping so the runtime can resolve
+        # types from lazily-loaded submodules by importing them on demand.
+        _reference_map._assembly_to_module[assembly.name] = assembly.module
+
         # Give our record of the assembly back to the caller.
         return assembly
 

packages/jsii-pacmak/lib/targets/python.ts

@@ -1797,6 +1797,9 @@ class PythonModule implements PythonType {
     code.line('import builtins');
     code.line('import datetime');
     code.line('import enum');
+    if (this.modules.length > 0) {
+      code.line('import importlib as _importlib');
+    }
     code.line('import typing');
     code.line();
     code.line('import jsii');
@@ -1920,29 +1923,88 @@ class PythonModule implements PythonType {
       code.line('__all__: typing.List[typing.Any] = []');
     }
 
+    // Emit TYPE_CHECKING block with explicit submodule imports so that
+    // static type checkers (pyright, mypy) can see the submodule names
+    // that are listed in __all__. At runtime TYPE_CHECKING is False,
+    // so these imports don't execute and lazy loading is preserved.
+    if (this.modules.length > 0) {
+      const submoduleNames = this.modules
+        .sort((l, r) => l.pythonName.localeCompare(r.pythonName))
+        .map((module) =>
+          module.pythonName.substring(this.pythonName.length + 1),
+        );
+
+      code.line();
+      code.line(
+        '# Type-checking-only imports for static analyzers (pyright/mypy).',
+      );
+      code.line(
+        '# At runtime TYPE_CHECKING is False, preserving lazy loading.',
+      );
+      code.openBlock('if typing.TYPE_CHECKING');
+      for (const name of submoduleNames) {
+        code.line(`from . import ${name} as ${name}`);
+      }
+      code.closeBlock();
+    }
+
     // Next up, we'll use publication to ensure that all of the non-public names
     // get hidden from dir(), tab-complete, etc.
     code.line();
     code.line('publication.publish()');
 
-    // Finally, we'll load all registered python modules
+    // Finally, we'll set up lazy loading for all registered python modules.
+    // We define __getattr__ and __dir__ and then install them on the public
+    // module (the one publication.publish() placed in sys.modules) so that
+    // lazy attribute access works through the publication barrier.
     if (this.modules.length > 0) {
       code.line();
-      code.line(
-        '# Loading modules to ensure their types are registered with the jsii runtime library',
-      );
-      for (const module of this.modules.sort((l, r) =>
-        l.pythonName.localeCompare(r.pythonName),
-      )) {
-        // Rather than generating an absolute import like
-        // "import jsii_calc.submodule" this builds a relative import like
-        // "from . import submodule". This enables distributing python packages
-        // and using the generated modules in the same codebase.
-        const submodule = module.pythonName.substring(
-          this.pythonName.length + 1,
+      // Build sorted list of submodule short names
+      const submoduleNames = this.modules
+        .sort((l, r) => l.pythonName.localeCompare(r.pythonName))
+        .map((module) =>
+          module.pythonName.substring(this.pythonName.length + 1),
         );
-        code.line(`from . import ${submodule}`);
+
+      // Emit _SUBMODULES set
+      code.indent('_SUBMODULES = {');
+      for (const name of submoduleNames) {
+        code.line(`"${name}",`);
       }
+      code.unindent('}');
+      code.line();
+
+      // Emit __getattr__ function
+      code.openBlock('def __getattr__(name: str) -> object');
+      code.openBlock('if name in _SUBMODULES');
+      code.line('mod = _importlib.import_module(f".{name}", __name__)');
+      code.line('globals()[name] = mod');
+      code.line('return mod');
+      code.closeBlock();
+      code.line(
+        'raise AttributeError(f"module {__name__!r} has no attribute {name!r}")',
+      );
+      code.closeBlock();
+      code.line();
+
+      // Emit __dir__ function — quoted return type because pyright flags
+      // bare `list[str]` as a runtime subscript error when pythonVersion < 3.9
+      code.openBlock('def __dir__() -> "list[str]"');
+      code.line('return [*__all__, *_SUBMODULES]');
+      code.closeBlock();
+      code.line();
+
+      // Install __getattr__ and __dir__ on the public module that
+      // publication.publish() placed in sys.modules. publication replaces
+      // the module object but doesn't copy __getattr__/__dir__, so without
+      // this, attribute access like `pkg.submodule` would raise AttributeError.
+      //
+      // We use setattr() instead of direct assignment because mypy treats
+      // __getattr__ and __dir__ as special methods on ModuleType and rejects
+      // direct assignment with "Cannot assign to a method [method-assign]".
+      code.line('import sys as _sys');
+      code.line('setattr(_sys.modules[__name__], "__getattr__", __getattr__)');
+      code.line('setattr(_sys.modules[__name__], "__dir__", __dir__)');
     }
 
     context.typeCheckingHelper.flushStubs(code);