|
| 1 | +"""Lazy module attributes (PEP 562) shared by the SDK's package inits. |
| 2 | +
|
| 3 | +`import mcp` no longer walks the SDK's whole module graph; each package instead |
| 4 | +resolves what it exports the first time it is asked for it. `lazy_module_attrs` |
| 5 | +builds the `__getattr__`/`__dir__` pair for such a module from two kinds of |
| 6 | +lazy names: |
| 7 | +
|
| 8 | +* `exports` - a re-exported object, described as `("home.module", "attr")` |
| 9 | + or as a zero-argument loader, imported/loaded on first access and then |
| 10 | + cached in the module namespace, so every later lookup is a plain attribute |
| 11 | + read; |
| 12 | +* `submodules` - the package's own submodules, imported when the attribute |
| 13 | + chain reaches them (`import mcp` then `mcp.client.stdio` used to work as a |
| 14 | + side effect of eager imports; this keeps such chains resolving). Pass the |
| 15 | + set explicitly, or `None` to use the package's real submodules (listed |
| 16 | + once, on first need). |
| 17 | +
|
| 18 | +A name that is neither raises `AttributeError` immediately: no filesystem |
| 19 | +search and no speculative import on a miss. Callers bind the pair under |
| 20 | +`if not TYPE_CHECKING:` so a static type checker never sees a module-level |
| 21 | +`__getattr__` (it would type every misspelled attribute as `object` instead of |
| 22 | +reporting it) while the mirrored `TYPE_CHECKING` imports keep the real names |
| 23 | +typed. |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import sys |
| 29 | +from collections.abc import Callable, Mapping |
| 30 | +from importlib import import_module |
| 31 | + |
| 32 | +__all__ = ["lazy_module_attrs"] |
| 33 | + |
| 34 | +Loader = tuple[str, str] | Callable[[], object] |
| 35 | +"""How a lazy export is produced: `("home.module", "attr")` or a zero-arg loader.""" |
| 36 | + |
| 37 | +_HELPER_NAMES = frozenset({"__getattr__", "__dir__"}) |
| 38 | +"""The two functions this helper adds to a module; never reported by its `__dir__`.""" |
| 39 | + |
| 40 | +DEFAULT_HIDDEN = frozenset({"TYPE_CHECKING", "_lazy_module_attrs"}) |
| 41 | +"""The conventional machinery names a caller binds to install the pair (hidden from `dir()`).""" |
| 42 | + |
| 43 | + |
| 44 | +def lazy_module_attrs( |
| 45 | + module_name: str, |
| 46 | + module_globals: dict[str, object], |
| 47 | + *, |
| 48 | + exports: Mapping[str, Loader] | None = None, |
| 49 | + submodules: frozenset[str] | None = None, |
| 50 | + hidden: frozenset[str] = DEFAULT_HIDDEN, |
| 51 | +) -> tuple[Callable[[str], object], Callable[[], list[str]]]: |
| 52 | + """Build a module's `__getattr__` and `__dir__` for lazily-resolved names. |
| 53 | +
|
| 54 | + Args: |
| 55 | + module_name: The `__name__` of the module the pair is installed on. |
| 56 | + module_globals: That module's `globals()`; resolved exports are cached |
| 57 | + in it, so `__getattr__` runs at most once per name. |
| 58 | + exports: Lazy re-exports, name -> `("home.module", "attr")` or a |
| 59 | + zero-argument loader. Resolved on first access. |
| 60 | + submodules: Submodule names reachable as attributes without an |
| 61 | + explicit import. `None` (the default) means the package's real |
| 62 | + submodules, discovered once on first need; a frozenset (possibly |
| 63 | + empty) means exactly those. |
| 64 | + hidden: Namespace entries that are lazy-loading machinery and must not |
| 65 | + appear in `dir(module)` (defaults to the two conventional binding |
| 66 | + names, `TYPE_CHECKING` and `_lazy_module_attrs`). |
| 67 | +
|
| 68 | + Returns: |
| 69 | + The `(__getattr__, __dir__)` pair to bind in the module namespace. |
| 70 | + """ |
| 71 | + export_map: Mapping[str, Loader] = exports or {} |
| 72 | + known_submodules = submodules |
| 73 | + |
| 74 | + def submodule_names() -> frozenset[str]: |
| 75 | + # An explicit set, or the real submodules of the package (one directory |
| 76 | + # listing, cached; dunder entries such as `__main__` are not attributes). |
| 77 | + # A plain module (no `__path__`) has no submodules. |
| 78 | + nonlocal known_submodules |
| 79 | + if known_submodules is None: |
| 80 | + # pkgutil is imported for the one listing rather than at module top, |
| 81 | + # keeping this helper (and so a bare `import mcp`) as light as possible. |
| 82 | + import pkgutil |
| 83 | + |
| 84 | + search_path = getattr(sys.modules[module_name], "__path__", None) |
| 85 | + known_submodules = frozenset( |
| 86 | + info.name |
| 87 | + for info in (pkgutil.iter_modules(search_path) if search_path else ()) |
| 88 | + if not info.name.startswith("__") |
| 89 | + ) |
| 90 | + return known_submodules |
| 91 | + |
| 92 | + def __getattr__(name: str) -> object: |
| 93 | + loader = export_map.get(name) |
| 94 | + if loader is not None: |
| 95 | + if callable(loader): |
| 96 | + value = loader() |
| 97 | + else: |
| 98 | + home_module, attr = loader |
| 99 | + value = getattr(import_module(home_module), attr) |
| 100 | + elif not name.startswith("__") and name in submodule_names(): |
| 101 | + # The import system also binds the submodule on the package, so |
| 102 | + # this too runs at most once per name. |
| 103 | + value = import_module(f"{module_name}.{name}") |
| 104 | + else: |
| 105 | + raise AttributeError(f"module {module_name!r} has no attribute {name!r}") |
| 106 | + module_globals[name] = value # cache: later accesses never reach __getattr__ |
| 107 | + return value |
| 108 | + |
| 109 | + def __dir__() -> list[str]: |
| 110 | + return sorted((module_globals.keys() - hidden - _HELPER_NAMES) | export_map.keys() | submodule_names()) |
| 111 | + |
| 112 | + return __getattr__, __dir__ |
0 commit comments