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

PiperOrigin-RevId: 838839559

Author: sararob

tests/unit/architecture/test_vertexai_import.py

@@ -81,17 +81,7 @@ 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) == sorted(
-        [
-            vertexai_module_name,
-            f"{vertexai_module_name}.types",
-        ]
-    )
-
-    placeholder_vertexai_types_module = sys.modules[f"{vertexai_module_name}.types"]
-    assert isinstance(placeholder_vertexai_types_module, vertexai._LazyTypesLoader)
-
-    assert f"{vertexai_module_name}._genai.types" not in sys.modules
+    assert sorted(new_modules_after_vertexai) == [vertexai_module_name]
 
     assert vertexai_import_timedelta.total_seconds() < 0.005
     assert aip_import_timedelta.total_seconds() < 40

vertexai/__init__.py

@@ -17,9 +17,6 @@
 import importlib
 import sys
 
-from types import ModuleType
-from typing import Any
-
 from google.cloud.aiplatform import version as aiplatform_version
 
 __version__ = aiplatform_version.__version__
@@ -30,23 +27,6 @@
 _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: Any, **kwargs: Any):
-        super().__init__(*args, **kwargs)
-        self._module = None
-
-    def __getattr__(self, name: str):
-        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/
@@ -65,7 +45,12 @@ def __getattr__(name):  # type: ignore[no-untyped-def]
         return getattr(_genai_client, name)
 
     if name == "types":
-        return sys.modules[__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
 
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")