Simplify the pip wheel info

- Return a dict with consistent fields
- Remove caching
- Remove type annotations
- Leverage the known wheel package dir value to calculate full paths

Author: AA-Turner

Lib/ensurepip/__init__.py

@@ -1,23 +1,15 @@
-import collections
 import os
 import os.path
 import subprocess
 import sys
 import sysconfig
 import tempfile
-from contextlib import suppress
-from functools import cache
 from importlib import resources
 
 
 __all__ = ["version", "bootstrap"]
 _PIP_VERSION = "23.3.2"
 
-# Packages bundled in ensurepip._bundled have wheel_name set.
-# Packages from WHEEL_PKG_DIR have wheel_path set.
-_Package = collections.namedtuple('Package',
-                                  ('version', 'wheel_name', 'wheel_path'))
-
 # Directory of system wheel packages. Some Linux distribution packaging
 # policies recommend against bundling dependencies. For example, Fedora
 # installs wheel packages in the /usr/share/python-wheels/ directory and don't
@@ -44,22 +36,17 @@ def _find_wheel_pkg_dir_pip():
 
         # Extract '21.2.4' from 'pip-21.2.4-py3-none-any.whl'
         version = filename.removeprefix("pip-").partition("-")[0]
-        wheel_path = os.path.join(_WHEEL_PKG_DIR, filename)
-        return _Package(version, None, wheel_path)
+        return {"version": version, "filename": filename, "bundled": False}
 
     return None
 
 
-@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)
-
-    with suppress(LookupError):
-        # only use the wheel package directory if pip wheel is found there
-        pip_pkg = _find_wheel_pkg_dir_pip(_WHEEL_PKG_DIR)
-
-    return pip_pkg
+def _get_pip_info():
+    # Prefer pip from the wheel package directory, if present.
+    if (pip_info := _find_wheel_pkg_dir_pip()) is not None:
+        return pip_info
+    filename = f"pip-{_PIP_VERSION}-py3-none-any.whl"
+    return {"version": _PIP_VERSION, "filename": filename, "bundled": True}
 
 
 def _run_pip(args, additional_paths=None):
@@ -92,7 +79,7 @@ def version():
     """
     Returns a string specifying the bundled version of pip.
     """
-    return _get_usable_pip_package().version
+    return _get_pip_info()["version"]
 
 
 def _disable_pip_configuration_settings():
@@ -154,17 +141,16 @@ 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
-        package = _get_usable_pip_package()
-        if package.wheel_name:
+        package = _get_pip_info()
+        wheel_name = package["filename"]
+        if package["bundled"]:
             # 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:
+            with open(os.path.join(_WHEEL_PKG_DIR, wheel_name), "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:
@@ -183,6 +169,7 @@ def _bootstrap(*, root=None, upgrade=False, user=False,
 
         return _run_pip([*args, "pip"], [filename])
 
+
 def _uninstall_helper(*, verbosity=0):
     """Helper to support a clean default uninstall process on Windows
 

Lib/test/test_ensurepip.py

@@ -6,19 +6,12 @@
 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()
@@ -27,43 +20,39 @@ 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, '_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
+    def test_get_pip_info_no_dir(self):
+        # Test _get_pip_info() without a wheel package directory
         with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', None):
-            pip_pkg = ensurepip._get_usable_pip_package()
+            pip_info = ensurepip._get_pip_info()
 
-            # when bundled pip wheel package is used, we get _PIP_VERSION
+            # when the bundled pip wheel is used, we get _PIP_VERSION
             self.assertEqual(ensurepip._PIP_VERSION, ensurepip.version())
 
-        # use bundled pip wheel package
-        self.assertIsNotNone(pip_pkg.wheel_name)
+        # use the bundled pip wheel
+        pip_filename = f'pip-{ensurepip._PIP_VERSION}-py3-none-any.whl'
+        expected = {"version": ensurepip._PIP_VERSION, "filename": pip_filename,
+                    "bundled": True}
+        self.assertDictEqual(pip_info, expected)
 
-    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"
+    def test_get_pip_info_with_dir(self):
+        # Test _get_pip_info() with a wheel package directory
         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
+            # not used, make sure that they're ignored
+            self.touch(tmpdir, "pip-1.2.3-py2.py3-none-any.whl")
             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")
+            self.touch(tmpdir, "pip-script.py")
 
-            with unittest.mock.patch.object(
-                    ensurepip, '_WHEEL_PKG_DIR', tmpdir,
-            ):
-                pip_pkg = ensurepip._get_usable_pip_package()
+            with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir):
+                pip_info = ensurepip._get_pip_info()
 
-            self.assertEqual(pip_pkg.version, '20.2.2')
-            self.assertEqual(pip_pkg.wheel_path,
-                             os.path.join(tmpdir, pip_filename))
+            expected = {"version": '20.2.2', "filename": pip_filename, "bundled": False}
+            self.assertDictEqual(pip_info, expected)
 
 
 class EnsurepipMixin:
@@ -102,30 +91,6 @@ 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/")