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]: ece20db [2]: #101039 [3]: #95299
python · pradyunsg · Jan 30, 2024 · Sep 10, 2023 · Sep 15, 2023 · Sep 16, 2023
commit c32777f19ec035a29ca8612937b3d3e9c8cd8549
@@ -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.3.2"
-_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):

@@ -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-10-23-35-20.gh-issue-109245.yCVQbo.rst b/Misc/NEWS.d/next/Library/2023-09-10-23-35-20.gh-issue-109245.yCVQbo.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.