From 46285bf11623cfe6ac77011be11996388f67f1d0 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 1 Dec 2025 10:38:44 -0800 Subject: [PATCH] fix: Enable `from vertexai.types import TypeName` without needing to run `from vertexai import types` first PiperOrigin-RevId: 838839559 --- .../unit/architecture/test_vertexai_import.py | 12 +-------- vertexai/__init__.py | 27 +++++-------------- 2 files changed, 7 insertions(+), 32 deletions(-) diff --git a/tests/unit/architecture/test_vertexai_import.py b/tests/unit/architecture/test_vertexai_import.py index de7b1370de..2170c7fa9e 100644 --- a/tests/unit/architecture/test_vertexai_import.py +++ b/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 diff --git a/vertexai/__init__.py b/vertexai/__init__.py index dbe3e44436..4a7c313a84 100644 --- a/vertexai/__init__.py +++ b/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}'")