From f4a6cbe6530a5a01b2c5fb95f388dfcc8e8e331f Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 1 Dec 2025 06:50:15 -0800 Subject: [PATCH] fix: Enable `from vertexai.types import TypeName` without needing to run `from vertexai import types` first PiperOrigin-RevId: 838752902 --- .../unit/architecture/test_vertexai_import.py | 12 ++++++++- vertexai/__init__.py | 27 ++++++++++++++----- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/tests/unit/architecture/test_vertexai_import.py b/tests/unit/architecture/test_vertexai_import.py index 2170c7fa9e..de7b1370de 100644 --- a/tests/unit/architecture/test_vertexai_import.py +++ b/tests/unit/architecture/test_vertexai_import.py @@ -81,7 +81,17 @@ 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) == 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 vertexai_import_timedelta.total_seconds() < 0.005 assert aip_import_timedelta.total_seconds() < 40 diff --git a/vertexai/__init__.py b/vertexai/__init__.py index 4a7c313a84..dbe3e44436 100644 --- a/vertexai/__init__.py +++ b/vertexai/__init__.py @@ -17,6 +17,9 @@ 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__ @@ -27,6 +30,23 @@ _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/ @@ -45,12 +65,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}'")