perf: defer eager type/streaming imports at import openai

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

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__()

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))

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}