Implement __init__.py lazy loading for codegen'd objects

crates/re_types_builder/src/codegen/python.rs

@@ -442,16 +442,41 @@ fn write_init_file(
         "
             from __future__ import annotations
 
+            from typing import TYPE_CHECKING, Any
             ",
-        0,
+        1,
     );
-    for (module, names) in mods {
-        let names = names.join(", ");
-        code.push_text(&format!("from .{module} import {names}"), 1, 0);
-    }
+
     if !manifest.is_empty() {
-        code.push_unindented_text(format!("\n__all__ = [{manifest}]"), 0);
+        code.push_unindented_text("if TYPE_CHECKING:", 1);
+        for (module, names) in mods {
+            let names = names.join(", ");
+            code.push_text(&format!("from .{module} import {names}"), 1, 4);
+        }
+
+        code.push_unindented_text(format!("\n\n__all__ = [{manifest}]"), 1);
+    }
+
+    code.push_text("\nmodule_content: dict[str, str] = {", 1, 0);
+    for (module, names) in mods {
+        for name in names {
+            code.push_text(&format!("{name:?}: {module:?},"), 1, 4);
+        }
     }
+    code.push_unindented_text("}", 1);
+
+    code.push_unindented_text(
+        r#"
+            def __getattr__(name: str) -> Any:
+                from importlib import import_module
+                if name in module_content:
+                    module = import_module(f".{module_content[name]}", __name__)
+                    return getattr(module, name)
+                raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+            "#,
+        0,
+    );
+
     files_to_write.insert(path, code);
 }
 

rerun_py/rerun_sdk/rerun/__init__.py

@@ -2,7 +2,7 @@
 
 import functools
 import random
-from typing import Any, Callable, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
 from uuid import UUID
 
 import numpy as np
@@ -126,48 +126,51 @@
 from ._image import ImageEncoded, ImageFormat
 from ._log import AsComponents, ComponentBatchLike, IndicatorComponentBatch, log, log_components
 from .any_value import AnyValues
-from .archetypes import (
-    AnnotationContext,
-    Arrows3D,
-    Asset3D,
-    BarChart,
-    Boxes2D,
-    Boxes3D,
-    Clear,
-    DepthImage,
-    DisconnectedSpace,
-    Image,
-    LineStrips2D,
-    LineStrips3D,
-    Mesh3D,
-    Pinhole,
-    Points2D,
-    Points3D,
-    SegmentationImage,
-    Tensor,
-    TextDocument,
-    TextLog,
-    TimeSeriesScalar,
-    Transform3D,
-    ViewCoordinates,
-)
+
+if TYPE_CHECKING:
+    from .archetypes import (
+        AnnotationContext,
+        Arrows3D,
+        Asset3D,
+        BarChart,
+        Boxes2D,
+        Boxes3D,
+        Clear,
+        DepthImage,
+        DisconnectedSpace,
+        Image,
+        LineStrips2D,
+        LineStrips3D,
+        Mesh3D,
+        Pinhole,
+        Points2D,
+        Points3D,
+        SegmentationImage,
+        Tensor,
+        TextDocument,
+        TextLog,
+        TimeSeriesScalar,
+        Transform3D,
+        ViewCoordinates,
+    )
+    from .components import (
+        Material,
+        MediaType,
+        MeshProperties,
+        OutOfTreeTransform3D,
+        OutOfTreeTransform3DBatch,
+        TextLogLevel,
+    )
+    from .datatypes import (
+        Quaternion,
+        RotationAxisAngle,
+        Scale3D,
+        TensorData,
+        TranslationAndMat3x3,
+        TranslationRotationScale3D,
+    )
+
 from .archetypes.boxes2d_ext import Box2DFormat
-from .components import (
-    Material,
-    MediaType,
-    MeshProperties,
-    OutOfTreeTransform3D,
-    OutOfTreeTransform3DBatch,
-    TextLogLevel,
-)
-from .datatypes import (
-    Quaternion,
-    RotationAxisAngle,
-    Scale3D,
-    TensorData,
-    TranslationAndMat3x3,
-    TranslationRotationScale3D,
-)
 from .error_utils import set_strict_mode
 from .log_deprecated.annotation import AnnotationInfo, ClassDescription, log_annotation_context
 from .log_deprecated.arrow import log_arrow
@@ -213,6 +216,21 @@
 from . import experimental  # isort: skip
 
 
+def __getattr__(name: str) -> Any:
+    from . import archetypes, components, datatypes
+
+    if name in archetypes.module_content:
+        return getattr(archetypes, name)
+
+    if name in components.module_content:
+        return getattr(components, name)
+
+    if name in datatypes.module_content:
+        return getattr(datatypes, name)
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 # =====================================
 # UTILITIES