diff --git a/launch/__init__.py b/launch/__init__.py index 4bbcff12..0781b712 100644 --- a/launch/__init__.py +++ b/launch/__init__.py @@ -8,9 +8,9 @@ # pylint: disable=C0413 import warnings +from importlib.metadata import version from typing import Sequence -import pkg_resources import pydantic if pydantic.VERSION.startswith("2."): @@ -32,7 +32,7 @@ SyncEndpoint, ) -__version__ = pkg_resources.get_distribution("scale-launch").version +__version__ = version("scale-launch") __all__: Sequence[str] = [ "AsyncEndpoint", "AsyncEndpointBatchResponse", diff --git a/launch/errors.py b/launch/errors.py index 05615755..df8b35ed 100644 --- a/launch/errors.py +++ b/launch/errors.py @@ -1,6 +1,6 @@ -import pkg_resources +from importlib.metadata import version -api_client_version = pkg_resources.get_distribution("scale-launch").version +api_client_version = version("scale-launch") INFRA_FLAKE_MESSAGES = [ "downstream duration timeout", diff --git a/launch/find_packages.py b/launch/find_packages.py index 2cc06200..faa5ef26 100644 --- a/launch/find_packages.py +++ b/launch/find_packages.py @@ -18,12 +18,19 @@ import logging import os import pkgutil +import re import sys import types import zipfile import zipimport from typing import Dict + +def _canonicalize_name(name: str) -> str: + """PEP 503 name normalization, matching the keys pkg_resources used to produce.""" + return re.sub(r"[-_.]+", "-", name).lower() + + EPP_NO_ERROR = 0 EPP_PKG_NOT_EXIST = 1 EPP_PKG_VERSION_MISMATCH = 2 @@ -106,27 +113,7 @@ def __init__(self): self.setuptools_module_set = set() self.nonlocal_package_path = set() - import pkg_resources - - # yixu: this populates either self.pip_pkg_map or self.nonlocal_package_path - # pkg_resources.working_set is basically a snapshot of sys.path, i.e. the packages that - # are imported - for dist in pkg_resources.working_set: # pylint: disable=not-an-iterable - module_path = dist.module_path or dist.location - if not module_path: - # Skip if no module path was found for pkg distribution - continue - - if os.path.realpath(module_path) != os.getcwd(): - # add to nonlocal_package path only if it's not current directory - self.nonlocal_package_path.add(module_path) - - self.pip_pkg_map[dist._key] = dist._version - for mn in dist._get_metadata("top_level.txt"): - if dist._key != "setuptools": - self.pip_module_map.setdefault(mn, []).append((dist._key, dist._version)) - else: - self.setuptools_module_set.add(mn) + self._index_installed_distributions() # yixu: searched_modules is basically just pkgutil.iter_modules self.searched_modules = {} @@ -142,12 +129,62 @@ def __init__(self): is_local = self.is_local_path(path) self.searched_modules[m.name] = ModuleInfo(m.name, path, is_local, m.ispkg) + def _index_installed_distributions(self): + # yixu: this populates either self.pip_pkg_map or self.nonlocal_package_path + # distributions() is basically a snapshot of sys.path, i.e. the packages that + # are imported + from importlib.metadata import distributions + + for dist in distributions(): + name = dist.metadata["Name"] + if not name: + # Skip malformed distributions with no name in their metadata + continue + + # pkg_resources keyed distributions by a normalized name; importlib.metadata + # reports the raw metadata name, so canonicalize to keep the requirement + # names emitted by seek_pip_packages() stable. + key = _canonicalize_name(name) + if key in self.pip_pkg_map: + # distributions() also yields shadowed copies (e.g. a vendored tree later + # on sys.path). First one wins, since that is the one that actually gets + # imported, which pkg_resources.working_set did implicitly. + continue + + # locate_file("") resolves to the directory the distribution was installed + # into (the parent of its .dist-info/.egg-info). realpath to match the + # normalized paths pkg_resources reported, since is_local_path() compares + # these against the entries collected here. + module_path = os.path.realpath(str(dist.locate_file(""))) + if not module_path: + # Skip if no module path was found for pkg distribution + continue + + if module_path != os.getcwd(): + # add to nonlocal_package path only if it's not current directory + self.nonlocal_package_path.add(module_path) + + self.pip_pkg_map[key] = dist.version + self._index_top_level_modules(dist, key) + + def _index_top_level_modules(self, dist, key): + for mn in (dist.read_text("top_level.txt") or "").splitlines(): + if not mn: + continue + if key != "setuptools": + self.pip_module_map.setdefault(mn, []).append((key, dist.version)) + else: + self.setuptools_module_set.add(mn) + def verify_pkg(self, pkg_req): - if pkg_req.name not in self.pip_pkg_map: + # pip_pkg_map is keyed by canonical name, so normalize the requirement name + # too: "jaraco.text", "jaraco_text" and "jaraco-text" all name one package. + req_key = _canonicalize_name(pkg_req.name) + if req_key not in self.pip_pkg_map: # package does not exist in the current python session return EPP_PKG_NOT_EXIST - if self.pip_pkg_map[pkg_req.name] not in pkg_req.specifier: + if self.pip_pkg_map[req_key] not in pkg_req.specifier: # package version being used in the current python session does not meet # the specified package version requirement return EPP_PKG_VERSION_MISMATCH diff --git a/tests/test_find_packages.py b/tests/test_find_packages.py new file mode 100644 index 00000000..ea539f16 --- /dev/null +++ b/tests/test_find_packages.py @@ -0,0 +1,136 @@ +import os + +import pytest +from packaging.requirements import Requirement + +from launch.find_packages import ( + EPP_NO_ERROR, + EPP_PKG_NOT_EXIST, + EPP_PKG_VERSION_MISMATCH, + ModuleManager, +) + + +class FakeDistribution: + """Stands in for an importlib.metadata.Distribution.""" + + def __init__(self, name, version, location, top_level=None): + self.metadata = {"Name": name} + self.version = version + self._location = location + self._top_level = top_level + + def locate_file(self, path): + return os.path.join(self._location, path) + + def read_text(self, filename): + return self._top_level if filename == "top_level.txt" else None + + +@pytest.fixture +def index_distributions(mocker): + """Build a ModuleManager over a fixed set of distributions.""" + + def _index(dists): + mocker.patch("importlib.metadata.distributions", return_value=iter(dists)) + return ModuleManager() + + return _index + + +def test_distribution_names_are_canonicalized(index_distributions, tmp_path): + # importlib.metadata reports the raw Name field, so "Typing_Extensions" and + # "jaraco.text" have to be normalized to the names requirements refer to. + manager = index_distributions( + [ + FakeDistribution("Typing_Extensions", "4.16.0", str(tmp_path), "typing_extensions\n"), + FakeDistribution("jaraco.text", "4.0.0", str(tmp_path), "jaraco\n"), + ] + ) + + assert manager.pip_pkg_map == {"typing-extensions": "4.16.0", "jaraco-text": "4.0.0"} + assert manager.pip_module_map["typing_extensions"] == [("typing-extensions", "4.16.0")] + + +def test_verify_pkg_matches_alternate_name_spellings(index_distributions, tmp_path): + manager = index_distributions( + [FakeDistribution("typing_extensions", "4.16.0", str(tmp_path), "typing_extensions\n")] + ) + + assert manager.verify_pkg(Requirement("typing-extensions>=4.0")) == EPP_NO_ERROR + assert manager.verify_pkg(Requirement("Typing_Extensions>=4.0")) == EPP_NO_ERROR + assert manager.verify_pkg(Requirement("typing.extensions>=4.0")) == EPP_NO_ERROR + assert manager.verify_pkg(Requirement("typing-extensions<4.0")) == EPP_PKG_VERSION_MISMATCH + assert manager.verify_pkg(Requirement("not-installed>=1.0")) == EPP_PKG_NOT_EXIST + + +def test_first_distribution_on_sys_path_wins(index_distributions, tmp_path): + # distributions() also yields shadowed copies, e.g. a vendored tree later on + # sys.path. The earlier one is the one that actually gets imported. + installed = tmp_path / "site-packages" + installed.mkdir() + vendored = tmp_path / "vendored" + vendored.mkdir() + + manager = index_distributions( + [ + FakeDistribution("packaging", "26.3", str(installed), "packaging\n"), + FakeDistribution("packaging", "26.0", str(vendored), "packaging\n"), + ] + ) + + assert manager.pip_pkg_map["packaging"] == "26.3" + assert manager.pip_module_map["packaging"] == [("packaging", "26.3")] + assert str(vendored) not in manager.nonlocal_package_path + + +def test_distribution_paths_are_realpath_normalized(index_distributions, tmp_path): + # is_local_path() compares these paths by identity, so an unresolved symlink + # would stop a nonlocal package from being recognized as one. + installed = tmp_path / "real-site-packages" + installed.mkdir() + symlinked = tmp_path / "linked-site-packages" + symlinked.symlink_to(installed) + + manager = index_distributions([FakeDistribution("some-pkg", "1.0.0", str(symlinked), "some_pkg\n")]) + + assert os.path.realpath(str(installed)) in manager.nonlocal_package_path + assert str(symlinked) not in manager.nonlocal_package_path + + +def test_distributions_without_a_name_are_skipped(index_distributions, tmp_path): + manager = index_distributions( + [ + FakeDistribution(None, "1.0.0", str(tmp_path), "broken\n"), + FakeDistribution("good-pkg", "2.0.0", str(tmp_path), "good_pkg\n"), + ] + ) + + assert manager.pip_pkg_map == {"good-pkg": "2.0.0"} + assert "broken" not in manager.pip_module_map + + +def test_missing_or_blank_top_level_metadata_is_tolerated(index_distributions, tmp_path): + # Wheels are not required to ship top_level.txt, and the ones that do may end + # with a trailing newline. + manager = index_distributions( + [ + FakeDistribution("no-top-level", "1.0.0", str(tmp_path), None), + FakeDistribution("blank-lines", "2.0.0", str(tmp_path), "blank_lines\n\n"), + ] + ) + + assert manager.pip_pkg_map == {"no-top-level": "1.0.0", "blank-lines": "2.0.0"} + assert manager.pip_module_map == {"blank_lines": [("blank-lines", "2.0.0")]} + + +def test_setuptools_modules_are_tracked_separately(index_distributions, tmp_path): + manager = index_distributions( + [ + FakeDistribution("setuptools", "80.10.2", str(tmp_path), "setuptools\npkg_resources\n"), + FakeDistribution("requests", "2.32.0", str(tmp_path), "requests\n"), + ] + ) + + assert manager.setuptools_module_set == {"setuptools", "pkg_resources"} + assert manager.pip_module_map == {"requests": [("requests", "2.32.0")]}