From a2b5fb15093cd08ef127efc1036f0ee74d57e4c6 Mon Sep 17 00:00:00 2001 From: Sviatoslav Sydorenko Date: Mon, 11 Sep 2023 01:23:21 +0200 Subject: [PATCH] Get rid of the ``ensurepip`` infra for many wheels This is a refactoring change that aims to simplify ``ensurepip``. Before it, this module had legacy infrastructure that made an assumption that ``ensurepip`` would be provisioning more then just a single wheel. That assumption is no longer true since [[1]][[2]][[3]]. In this change, the improvement is done around removing unnecessary loops and supporting structures to change the assumptions to expect only the bundled or replacement ``pip`` wheel. [1]: https://github.com/python/cpython/commit/ece20db [2]: https://github.com/python/cpython/pull/101039 [2]: https://github.com/python/cpython/issues/95299 --- Lib/ensurepip/__init__.py | 100 +++++++++--------- Lib/test/test_ensurepip.py | 65 +++++++++--- .../next/Library/2023-09-11-01-20-16.rst | 3 + 3 files changed, 101 insertions(+), 67 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2023-09-11-01-20-16.rst diff --git a/Lib/ensurepip/__init__.py b/Lib/ensurepip/__init__.py index 1fb1d505cfd0c53..dda2d198a336894 100644 --- a/Lib/ensurepip/__init__.py +++ b/Lib/ensurepip/__init__.py @@ -5,15 +5,13 @@ import sys import sysconfig import tempfile +from contextlib import suppress +from functools import cache from importlib import resources __all__ = ["version", "bootstrap"] -_PACKAGE_NAMES = ('pip',) _PIP_VERSION = "23.2.1" -_PROJECTS = [ - ("pip", _PIP_VERSION, "py3"), -] # Packages bundled in ensurepip._bundled have wheel_name set. # Packages from WHEEL_PKG_DIR have wheel_path set. @@ -27,8 +25,13 @@ _WHEEL_PKG_DIR = sysconfig.get_config_var('WHEEL_PKG_DIR') -def _find_packages(path): - packages = {} +def _find_packages(path: str | None) -> _Package: + if path is None: + raise LookupError( + 'The compile-time `WHEEL_PKG_DIR` is unset so there is ' + 'no place for looking up the wheels.', + ) + try: filenames = os.listdir(path) except OSError: @@ -38,41 +41,39 @@ def _find_packages(path): # of the same package, but don't attempt to implement correct version # comparison since this case should not happen. filenames = sorted(filenames) + pip_pkg = None for filename in filenames: # filename is like 'pip-21.2.4-py3-none-any.whl' if not filename.endswith(".whl"): continue - for name in _PACKAGE_NAMES: - prefix = name + '-' - if filename.startswith(prefix): - break - else: + if not filename.startswith('pip-'): continue # Extract '21.2.4' from 'pip-21.2.4-py3-none-any.whl' - version = filename.removeprefix(prefix).partition('-')[0] + discovered_pip_pkg_version = filename.removeprefix( + 'pip-', + ).partition('-')[0] wheel_path = os.path.join(path, filename) - packages[name] = _Package(version, None, wheel_path) - return packages + pip_pkg = _Package(discovered_pip_pkg_version, None, wheel_path) + + if pip_pkg is None: + raise LookupError( + '`WHEEL_PKG_DIR` does not contain any wheel files for `pip`.', + ) + + return pip_pkg -def _get_packages(): - global _PACKAGES, _WHEEL_PKG_DIR - if _PACKAGES is not None: - return _PACKAGES +@cache +def _get_usable_pip_package() -> _Package: + wheel_name = f"pip-{_PIP_VERSION}-py3-none-any.whl" + pip_pkg = _Package(_PIP_VERSION, wheel_name, None) - packages = {} - for name, version, py_tag in _PROJECTS: - wheel_name = f"{name}-{version}-{py_tag}-none-any.whl" - packages[name] = _Package(version, wheel_name, None) - if _WHEEL_PKG_DIR: - dir_packages = _find_packages(_WHEEL_PKG_DIR) - # only used the wheel package directory if all packages are found there - if all(name in dir_packages for name in _PACKAGE_NAMES): - packages = dir_packages - _PACKAGES = packages - return packages -_PACKAGES = None + with suppress(LookupError): + # only use the wheel package directory if all packages are found there + pip_pkg = _find_packages(_WHEEL_PKG_DIR) + + return pip_pkg def _run_pip(args, additional_paths=None): @@ -105,7 +106,7 @@ def version(): """ Returns a string specifying the bundled version of pip. """ - return _get_packages()['pip'].version + return _get_usable_pip_package().version def _disable_pip_configuration_settings(): @@ -167,24 +168,21 @@ def _bootstrap(*, root=None, upgrade=False, user=False, with tempfile.TemporaryDirectory() as tmpdir: # Put our bundled wheels into a temporary directory and construct the # additional paths that need added to sys.path - additional_paths = [] - for name, package in _get_packages().items(): - if package.wheel_name: - # Use bundled wheel package - wheel_name = package.wheel_name - wheel_path = resources.files("ensurepip") / "_bundled" / wheel_name - whl = wheel_path.read_bytes() - else: - # Use the wheel package directory - with open(package.wheel_path, "rb") as fp: - whl = fp.read() - wheel_name = os.path.basename(package.wheel_path) - - filename = os.path.join(tmpdir, wheel_name) - with open(filename, "wb") as fp: - fp.write(whl) - - additional_paths.append(filename) + package = _get_usable_pip_package() + if package.wheel_name: + # Use bundled wheel package + wheel_name = package.wheel_name + wheel_path = resources.files("ensurepip") / "_bundled" / wheel_name + whl = wheel_path.read_bytes() + else: + # Use the wheel package directory + with open(package.wheel_path, "rb") as fp: + whl = fp.read() + wheel_name = os.path.basename(package.wheel_path) + + filename = os.path.join(tmpdir, wheel_name) + with open(filename, "wb") as fp: + fp.write(whl) # Construct the arguments to be passed to the pip command args = ["install", "--no-cache-dir", "--no-index", "--find-links", tmpdir] @@ -197,7 +195,7 @@ def _bootstrap(*, root=None, upgrade=False, user=False, if verbosity: args += ["-" + "v" * verbosity] - return _run_pip([*args, *_PACKAGE_NAMES], additional_paths) + return _run_pip([*args, "pip"], [filename]) def _uninstall_helper(*, verbosity=0): """Helper to support a clean default uninstall process on Windows @@ -227,7 +225,7 @@ def _uninstall_helper(*, verbosity=0): if verbosity: args += ["-" + "v" * verbosity] - return _run_pip([*args, *reversed(_PACKAGE_NAMES)]) + return _run_pip([*args, "pip"]) def _main(argv=None): diff --git a/Lib/test/test_ensurepip.py b/Lib/test/test_ensurepip.py index 69ab2a4feaa9389..e190025c54035a3 100644 --- a/Lib/test/test_ensurepip.py +++ b/Lib/test/test_ensurepip.py @@ -6,12 +6,19 @@ import test.support import unittest import unittest.mock +from pathlib import Path import ensurepip import ensurepip._uninstall class TestPackages(unittest.TestCase): + def setUp(self): + ensurepip._get_usable_pip_package.cache_clear() + + def tearDown(self): + ensurepip._get_usable_pip_package.cache_clear() + def touch(self, directory, filename): fullname = os.path.join(directory, filename) open(fullname, "wb").close() @@ -20,42 +27,44 @@ def test_version(self): # Test version() with tempfile.TemporaryDirectory() as tmpdir: self.touch(tmpdir, "pip-1.2.3b1-py2.py3-none-any.whl") - with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None), - unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir)): + with unittest.mock.patch.object( + ensurepip, '_WHEEL_PKG_DIR', tmpdir, + ): self.assertEqual(ensurepip.version(), '1.2.3b1') def test_get_packages_no_dir(self): # Test _get_packages() without a wheel package directory - with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None), - unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', None)): - packages = ensurepip._get_packages() + with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', None): + pip_pkg = ensurepip._get_usable_pip_package() - # when bundled wheel packages are used, we get _PIP_VERSION + # when bundled pip wheel package is used, we get _PIP_VERSION self.assertEqual(ensurepip._PIP_VERSION, ensurepip.version()) - # use bundled wheel packages - self.assertIsNotNone(packages['pip'].wheel_name) + # use bundled pip wheel package + self.assertIsNotNone(pip_pkg.wheel_name) def test_get_packages_with_dir(self): # Test _get_packages() with a wheel package directory + older_pip_filename = "pip-1.2.3-py2.py3-none-any.whl" pip_filename = "pip-20.2.2-py2.py3-none-any.whl" with tempfile.TemporaryDirectory() as tmpdir: + self.touch(tmpdir, older_pip_filename) self.touch(tmpdir, pip_filename) # not used, make sure that it's ignored self.touch(tmpdir, "wheel-0.34.2-py2.py3-none-any.whl") + # not used, make sure that it's ignored + self.touch(tmpdir, "non-whl") - with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None), - unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir)): - packages = ensurepip._get_packages() + with unittest.mock.patch.object( + ensurepip, '_WHEEL_PKG_DIR', tmpdir, + ): + pip_pkg = ensurepip._get_usable_pip_package() - self.assertEqual(packages['pip'].version, '20.2.2') - self.assertEqual(packages['pip'].wheel_path, + self.assertEqual(pip_pkg.version, '20.2.2') + self.assertEqual(pip_pkg.wheel_path, os.path.join(tmpdir, pip_filename)) - # wheel package is ignored - self.assertEqual(sorted(packages), ['pip']) - class EnsurepipMixin: @@ -93,6 +102,30 @@ def test_basic_bootstrapping(self): additional_paths = self.run_pip.call_args[0][1] self.assertEqual(len(additional_paths), 1) + + def test_replacement_wheel_bootstrapping(self): + ensurepip._get_usable_pip_package.cache_clear() + + pip_wheel_name = ( + f'pip-{ensurepip._PIP_VERSION !s}-' + 'py3-none-any.whl' + ) + + with tempfile.TemporaryDirectory() as tmpdir: + tmp_path = Path(tmpdir) + tmp_wheel_path = tmp_path / pip_wheel_name + tmp_wheel_path.touch() + + with unittest.mock.patch.object( + ensurepip, '_WHEEL_PKG_DIR', tmpdir, + ): + ensurepip.bootstrap() + + ensurepip._get_usable_pip_package.cache_clear() + + additional_paths = self.run_pip.call_args[0][1] + self.assertEqual(Path(additional_paths[-1]).name, pip_wheel_name) + def test_bootstrapping_with_root(self): ensurepip.bootstrap(root="/foo/bar/") diff --git a/Misc/NEWS.d/next/Library/2023-09-11-01-20-16.rst b/Misc/NEWS.d/next/Library/2023-09-11-01-20-16.rst new file mode 100644 index 000000000000000..41a56fb30b62fd1 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2023-09-11-01-20-16.rst @@ -0,0 +1,3 @@ +Simplified ``ensurepip`` to stop assuming that it can provision multiple +wheels. The refreshed implementation now expects to only provision +a ``pip`` wheel and no other distribution packages.