fix: Remove lazy loading of vertexai.types module to enable from vertexai.types import TypeName

PiperOrigin-RevId: 834811652

Author: sararob

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 actual types module when an attribute is accessed."""
+
+    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)
+
+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}'")