diff --git a/litellm/litellm_core_utils/deprecated_package.py b/litellm/litellm_core_utils/deprecated_package.py index 18cc665a6b9..542686ecf19 100644 --- a/litellm/litellm_core_utils/deprecated_package.py +++ b/litellm/litellm_core_utils/deprecated_package.py @@ -1,20 +1,53 @@ import importlib -import pkgutil import sys import warnings +from itertools import chain +from pathlib import Path from types import ModuleType +from typing import Final + + +def _dotted_name(root: Path, path: Path, package_name: str) -> str: + parts: Final = path.relative_to(root).with_suffix("").parts + return ".".join((package_name, *parts)).removesuffix(".__init__") + + +def _parents_and_self(dotted: str, package_name: str) -> tuple[str, ...]: + parts: Final = dotted.removeprefix(f"{package_name}.").split(".") + return tuple(f"{package_name}.{'.'.join(parts[:depth])}" for depth in range(1, len(parts) + 1)) + + +def _module_names(package: ModuleType) -> tuple[str, ...]: + roots: Final = tuple(Path(root) for root in package.__path__) + files: Final = tuple( + (root, path) for root in roots for path in root.rglob("*.py") + ) # comprehension-ok: nested walk over each package root + leaves: Final = tuple( + name + for name in (_dotted_name(root, path, package.__name__) for root, path in files) + if name != package.__name__ + ) + return tuple( + sorted( + set(chain.from_iterable(_parents_and_self(leaf, package.__name__) for leaf in leaves)) - {package.__name__} + ) + ) + + +def _alias(old_name: str, new_package_name: str, new_dotted: str) -> None: + module: Final = importlib.import_module(new_dotted) + old_dotted: Final = f"{old_name}{new_dotted.removeprefix(new_package_name)}" + old_parent, _, child = old_dotted.rpartition(".") + sys.modules[old_dotted] = module + setattr(sys.modules[old_parent], child, module) def alias_deprecated_package(old_name: str, new_package: ModuleType) -> None: - """Make every submodule of ``new_package`` importable under ``old_name`` and warn once.""" + """Make every module under ``new_package`` importable and attribute-reachable under ``old_name`` and warn once.""" warnings.warn( f"{old_name} is deprecated, import {new_package.__name__} instead", DeprecationWarning, stacklevel=3, ) - sys.modules.update( - { - f"{old_name}{info.name.removeprefix(new_package.__name__)}": importlib.import_module(info.name) - for info in pkgutil.walk_packages(new_package.__path__, prefix=f"{new_package.__name__}.") - } - ) + for new_dotted in _module_names(new_package): + _alias(old_name, new_package.__name__, new_dotted) diff --git a/tests/unit/experimental_mcp_client/test_init.py b/tests/unit/experimental_mcp_client/test_init.py index 8a1787f5ac5..9d2f1137470 100644 --- a/tests/unit/experimental_mcp_client/test_init.py +++ b/tests/unit/experimental_mcp_client/test_init.py @@ -1,5 +1,6 @@ import importlib import sys +from collections.abc import Mapping import pytest @@ -23,3 +24,6 @@ def test_old_package_warns_and_deep_imports_resolve_to_mcp_client(monkeypatch: p assert old.call_openai_tool is litellm.mcp_client.call_openai_tool assert importlib.import_module(f"{OLD_PACKAGE}.client") is litellm.mcp_client.client assert importlib.import_module(f"{OLD_PACKAGE}.tools") is litellm.mcp_client.tools + attrs: Mapping[str, object] = vars(old) + assert attrs["tools"] is litellm.mcp_client.tools + assert attrs["client"] is litellm.mcp_client.client diff --git a/tests/unit/proxy/_experimental/mcp_server/test_init.py b/tests/unit/proxy/_experimental/mcp_server/test_init.py index 25016868950..2f933e11351 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_init.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_init.py @@ -1,8 +1,10 @@ import importlib import sys +from collections.abc import Mapping import pytest +import litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import litellm.proxy.mcp_server.server import litellm.proxy.mcp_server.utils @@ -15,7 +17,14 @@ def test_old_package_warns_and_deep_imports_resolve_to_proxy_mcp_server(monkeypa monkeypatch.delitem(sys.modules, name) with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"): - importlib.import_module(OLD_PACKAGE) + old = importlib.import_module(OLD_PACKAGE) assert importlib.import_module(f"{OLD_PACKAGE}.server") is litellm.proxy.mcp_server.server assert importlib.import_module(f"{OLD_PACKAGE}.utils") is litellm.proxy.mcp_server.utils + assert ( + importlib.import_module(f"{OLD_PACKAGE}.auth.user_api_key_auth_mcp") + is litellm.proxy.mcp_server.auth.user_api_key_auth_mcp + ) + attrs: Mapping[str, object] = vars(old) + assert attrs["server"] is litellm.proxy.mcp_server.server + assert attrs["auth"] is litellm.proxy.mcp_server.auth