diff --git a/src/openai/__init__.py b/src/openai/__init__.py index a1d950d3b4..fa295b1d94 100644 --- a/src/openai/__init__.py +++ b/src/openai/__init__.py @@ -4,9 +4,9 @@ import os as _os import typing as _t +import importlib as _importlib from typing_extensions import override -from . import types from ._types import NOT_GIVEN, Omit, NoneType, NotGiven, Transport, ProxiesTypes, omit, not_given from ._utils import file_from_path from ._client import Client, OpenAI, Stream, Timeout, Transport, AsyncClient, AsyncOpenAI, AsyncStream, RequestOptions @@ -41,8 +41,6 @@ from ._utils._logs import setup_logging as _setup_logging from ._data_residency import DataResidency from ._legacy_response import HttpxBinaryResponseContent as HttpxBinaryResponseContent -from .types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides - __all__ = [ "types", "__version__", @@ -100,18 +98,46 @@ "WebSocketConnectionClosedError", ] -if not _t.TYPE_CHECKING: +if _t.TYPE_CHECKING: + from . import types as types + from .lib import pydantic_function_tool as pydantic_function_tool + from .lib.streaming import ( + AssistantEventHandler as AssistantEventHandler, + AsyncAssistantEventHandler as AsyncAssistantEventHandler, + ) + from .types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides +else: from ._utils._resources_proxy import resources as resources + from ._utils._types_proxy import types as types -from .lib import azure as _azure, bedrock as _bedrock, pydantic_function_tool as pydantic_function_tool +from .lib import azure as _azure, bedrock as _bedrock from .version import VERSION as VERSION from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI from .lib.bedrock import BedrockOpenAI as BedrockOpenAI, AsyncBedrockOpenAI as AsyncBedrockOpenAI from .lib._old_api import * -from .lib.streaming import ( - AssistantEventHandler as AssistantEventHandler, - AsyncAssistantEventHandler as AsyncAssistantEventHandler, -) + +_LAZY_IMPORTS = { + "pydantic_function_tool": (".lib", "pydantic_function_tool"), + "AssistantEventHandler": (".lib.streaming", "AssistantEventHandler"), + "AsyncAssistantEventHandler": (".lib.streaming", "AsyncAssistantEventHandler"), + "ReconnectingEvent": (".types.websocket_reconnection", "ReconnectingEvent"), + "ReconnectingOverrides": (".types.websocket_reconnection", "ReconnectingOverrides"), +} + + +def __getattr__(name: str) -> _t.Any: + lazy_import = _LAZY_IMPORTS.get(name) + if lazy_import is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module_name, attribute_name = lazy_import + value = getattr(_importlib.import_module(module_name, __name__), attribute_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(_LAZY_IMPORTS)) _setup_logging() @@ -123,7 +149,9 @@ for __name in __all__: if not __name.startswith("__"): try: - __locals[__name].__module__ = "openai" + __object = __locals.get(__name) + if __object is not None: + __object.__module__ = "openai" except (TypeError, AttributeError): # Some of our exported symbols are builtins which we can't set attributes for. pass diff --git a/src/openai/_utils/_types_proxy.py b/src/openai/_utils/_types_proxy.py new file mode 100644 index 0000000000..54ea634b76 --- /dev/null +++ b/src/openai/_utils/_types_proxy.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from typing import Any +from typing_extensions import override + +from ._proxy import LazyProxy + + +class TypesProxy(LazyProxy[Any]): + """Lazily import ``openai.types`` when an exported type is accessed.""" + + @override + def __load__(self) -> Any: + import importlib + + return importlib.import_module("openai.types") + + +types = TypesProxy().__as_proxied__() diff --git a/src/openai/lib/__init__.py b/src/openai/lib/__init__.py index 5c6cb782c0..9428b14e17 100644 --- a/src/openai/lib/__init__.py +++ b/src/openai/lib/__init__.py @@ -1,2 +1,30 @@ -from ._tools import pydantic_function_tool as pydantic_function_tool -from ._parsing import ResponseFormatT as ResponseFormatT +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from ._tools import pydantic_function_tool as pydantic_function_tool + from ._parsing import ResponseFormatT as ResponseFormatT + +__all__ = ["ResponseFormatT", "pydantic_function_tool"] + +_LAZY_IMPORTS = { + "ResponseFormatT": ("._parsing", "ResponseFormatT"), + "pydantic_function_tool": ("._tools", "pydantic_function_tool"), +} + + +def __getattr__(name: str) -> Any: + lazy_import = _LAZY_IMPORTS.get(name) + if lazy_import is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module_name, attribute_name = lazy_import + value = getattr(importlib.import_module(module_name, __name__), attribute_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(_LAZY_IMPORTS)) diff --git a/tests/test_lazy_imports.py b/tests/test_lazy_imports.py new file mode 100644 index 0000000000..187eca9f74 --- /dev/null +++ b/tests/test_lazy_imports.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import openai + + +def _run_fresh_import(source: str) -> dict[str, object]: + env = os.environ.copy() + env["PYTHONPATH"] = str(Path(__file__).parents[1] / "src") + result = subprocess.run( + [sys.executable, "-c", source], + check=True, + capture_output=True, + text=True, + env=env, + ) + return json.loads(result.stdout) + + +def test_top_level_import_defers_optional_subtrees() -> None: + # `import openai` should no longer eagerly build the Assistants ``beta`` type + # tree (the largest unused subtree) or the streaming helpers. They must load + # only on first use. + imported = _run_fresh_import( + """ +import json +import sys +import openai + +print(json.dumps({ + "beta": "openai.types.beta" in sys.modules, + "streaming": "openai.lib.streaming" in sys.modules, +})) +""" + ) + + assert imported == {"beta": False, "streaming": False} + + +def test_lazy_top_level_exports_preserve_public_objects() -> None: + from openai.lib import pydantic_function_tool + from openai.lib.streaming import AssistantEventHandler, AsyncAssistantEventHandler + + assert openai.pydantic_function_tool is pydantic_function_tool + assert openai.AssistantEventHandler is AssistantEventHandler + assert openai.AsyncAssistantEventHandler is AsyncAssistantEventHandler + + +def test_types_proxy_loads_real_module_on_access() -> None: + imported = _run_fresh_import( + """ +import json +import sys +import openai + +batch = openai.types.Batch +print(json.dumps({ + "name": batch.__name__, + "types": "openai.types" in sys.modules, +})) +""" + ) + + assert imported == {"name": "Batch", "types": True}