fix: Enable from vertexai.types import TypeName without needing to run from vertexai import types first

PiperOrigin-RevId: 834811652

Author: sararob

tests/unit/architecture/test_vertexai_import.py

@@ -81,7 +81,9 @@ def test_vertexai_import():
     new_modules_after_vertexai = modules_after_vertexai - modules_before_vertexai
 
     vertexai_module_name = vertexai.__name__  # == "vertexai"
-    assert sorted(new_modules_after_vertexai) == [vertexai_module_name]
+    assert sorted(new_modules_after_vertexai) == [
+        vertexai_module_name + "google.cloud.aiplatform.vertexai.types"
+    ]
 
     assert vertexai_import_timedelta.total_seconds() < 0.005
     assert aip_import_timedelta.total_seconds() < 40

vertexai/__init__.py

@@ -17,6 +17,8 @@
 import importlib
 import sys
 
+from types import ModuleType
+
 from google.cloud.aiplatform import version as aiplatform_version
 
 __version__ = aiplatform_version.__version__
@@ -27,6 +29,22 @@
 _genai_types = None
 
 
+class _LazyTypesLoader(ModuleType):
+    """A module that lazily loads the types module when an attribute is accessed via from `vertexai.types import TypeName`."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._module = None
+
+    def __getattr__(self, name):
+        if self._module is None:
+            self._module = importlib.import_module("._genai.types", __package__)
+        return getattr(self._module, name)
+
+# Register a placeholder _LazyTypesLoader instance for vertexai.types until it is accessed.
+sys.modules[__name__ + ".types"] = _LazyTypesLoader(__name__ + ".types")
+
+
 def __getattr__(name):  # type: ignore[no-untyped-def]
     # Lazy importing the preview submodule
     # See https://peps.python.org/pep-0562/
@@ -45,12 +63,7 @@ def __getattr__(name):  # type: ignore[no-untyped-def]
         return getattr(_genai_client, name)
 
     if name == "types":
-        global _genai_types
-        if _genai_types is None:
-            _genai_types = importlib.import_module("._genai.types", __name__)
-        if "vertexai.types" not in sys.modules:
-            sys.modules["vertexai.types"] = _genai_types
-        return _genai_types
+        return sys.modules[__name__ + ".types"]
 
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")