Merge pull request #3794 from youknowone/ensure-pip

Add ensurepip

Lib/ensurepip/__init__.py

@@ -0,0 +1,284 @@
+import collections
+import os
+import os.path
+import subprocess
+import sys
+import sysconfig
+import tempfile
+from importlib import resources
+
+
+
+__all__ = ["version", "bootstrap"]
+_PACKAGE_NAMES = ('setuptools', 'pip')
+_SETUPTOOLS_VERSION = "58.1.0"
+_PIP_VERSION = "22.0.4"
+_PROJECTS = [
+    ("setuptools", _SETUPTOOLS_VERSION, "py3"),
+    ("pip", _PIP_VERSION, "py3"),
+]
+
+# 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
+# install the ensurepip._bundled package.
+_WHEEL_PKG_DIR = sysconfig.get_config_var('WHEEL_PKG_DIR')
+
+
+def _find_packages(path):
+    packages = {}
+    try:
+        filenames = os.listdir(path)
+    except OSError:
+        # Ignore: path doesn't exist or permission error
+        filenames = ()
+    # Make the code deterministic if a directory contains multiple wheel files
+    # of the same package, but don't attempt to implement correct version
+    # comparison since this case should not happen.
+    filenames = sorted(filenames)
+    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:
+            continue
+
+        # Extract '21.2.4' from 'pip-21.2.4-py3-none-any.whl'
+        version = filename.removeprefix(prefix).partition('-')[0]
+        wheel_path = os.path.join(path, filename)
+        packages[name] = _Package(version, None, wheel_path)
+    return packages
+
+
+def _get_packages():
+    global _PACKAGES, _WHEEL_PKG_DIR
+    if _PACKAGES is not None:
+        return _PACKAGES
+
+    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
+
+
+def _run_pip(args, additional_paths=None):
+    # Run the bootstraping in a subprocess to avoid leaking any state that happens
+    # after pip has executed. Particulary, this avoids the case when pip holds onto
+    # the files in *additional_paths*, preventing us to remove them at the end of the
+    # invocation.
+    code = f"""
+import runpy
+import sys
+sys.path = {additional_paths or []} + sys.path
+sys.argv[1:] = {args}
+runpy.run_module("pip", run_name="__main__", alter_sys=True)
+"""
+    return subprocess.run([sys.executable, '-W', 'ignore::DeprecationWarning',
+                           "-c", code], check=True).returncode
+
+
+def version():
+    """
+    Returns a string specifying the bundled version of pip.
+    """
+    return _get_packages()['pip'].version
+
+
+def _disable_pip_configuration_settings():
+    # We deliberately ignore all pip environment variables
+    # when invoking pip
+    # See http://bugs.python.org/issue19734 for details
+    keys_to_remove = [k for k in os.environ if k.startswith("PIP_")]
+    for k in keys_to_remove:
+        del os.environ[k]
+    # We also ignore the settings in the default pip configuration file
+    # See http://bugs.python.org/issue20053 for details
+    os.environ['PIP_CONFIG_FILE'] = os.devnull
+
+
+def bootstrap(*, root=None, upgrade=False, user=False,
+              altinstall=False, default_pip=False,
+              verbosity=0):
+    """
+    Bootstrap pip into the current Python installation (or the given root
+    directory).
+
+    Note that calling this function will alter both sys.path and os.environ.
+    """
+    # Discard the return value
+    _bootstrap(root=root, upgrade=upgrade, user=user,
+               altinstall=altinstall, default_pip=default_pip,
+               verbosity=verbosity)
+
+
+def _bootstrap(*, root=None, upgrade=False, user=False,
+              altinstall=False, default_pip=False,
+              verbosity=0):
+    """
+    Bootstrap pip into the current Python installation (or the given root
+    directory). Returns pip command status code.
+
+    Note that calling this function will alter both sys.path and os.environ.
+    """
+    if altinstall and default_pip:
+        raise ValueError("Cannot use altinstall and default_pip together")
+
+    sys.audit("ensurepip.bootstrap", root)
+
+    _disable_pip_configuration_settings()
+
+    # By default, installing pip and setuptools installs all of the
+    # following scripts (X.Y == running Python version):
+    #
+    #   pip, pipX, pipX.Y, easy_install, easy_install-X.Y
+    #
+    # pip 1.5+ allows ensurepip to request that some of those be left out
+    if altinstall:
+        # omit pip, pipX and easy_install
+        os.environ["ENSUREPIP_OPTIONS"] = "altinstall"
+    elif not default_pip:
+        # omit pip and easy_install
+        os.environ["ENSUREPIP_OPTIONS"] = "install"
+
+    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
+                from ensurepip import _bundled
+                wheel_name = package.wheel_name
+                whl = resources.read_binary(_bundled, wheel_name)
+            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)
+
+        # Construct the arguments to be passed to the pip command
+        args = ["install", "--no-cache-dir", "--no-index", "--find-links", tmpdir]
+        if root:
+            args += ["--root", root]
+        if upgrade:
+            args += ["--upgrade"]
+        if user:
+            args += ["--user"]
+        if verbosity:
+            args += ["-" + "v" * verbosity]
+
+        return _run_pip([*args, *_PACKAGE_NAMES], additional_paths)
+
+def _uninstall_helper(*, verbosity=0):
+    """Helper to support a clean default uninstall process on Windows
+
+    Note that calling this function may alter os.environ.
+    """
+    # Nothing to do if pip was never installed, or has been removed
+    try:
+        import pip
+    except ImportError:
+        return
+
+    # If the installed pip version doesn't match the available one,
+    # leave it alone
+    available_version = version()
+    if pip.__version__ != available_version:
+        print(f"ensurepip will only uninstall a matching version "
+              f"({pip.__version__!r} installed, "
+              f"{available_version!r} available)",
+              file=sys.stderr)
+        return
+
+    _disable_pip_configuration_settings()
+
+    # Construct the arguments to be passed to the pip command
+    args = ["uninstall", "-y", "--disable-pip-version-check"]
+    if verbosity:
+        args += ["-" + "v" * verbosity]
+
+    return _run_pip([*args, *reversed(_PACKAGE_NAMES)])
+
+
+def _main(argv=None):
+    import argparse
+    parser = argparse.ArgumentParser(prog="python -m ensurepip")
+    parser.add_argument(
+        "--version",
+        action="version",
+        version="pip {}".format(version()),
+        help="Show the version of pip that is bundled with this Python.",
+    )
+    parser.add_argument(
+        "-v", "--verbose",
+        action="count",
+        default=0,
+        dest="verbosity",
+        help=("Give more output. Option is additive, and can be used up to 3 "
+              "times."),
+    )
+    parser.add_argument(
+        "-U", "--upgrade",
+        action="store_true",
+        default=False,
+        help="Upgrade pip and dependencies, even if already installed.",
+    )
+    parser.add_argument(
+        "--user",
+        action="store_true",
+        default=False,
+        help="Install using the user scheme.",
+    )
+    parser.add_argument(
+        "--root",
+        default=None,
+        help="Install everything relative to this alternate root directory.",
+    )
+    parser.add_argument(
+        "--altinstall",
+        action="store_true",
+        default=False,
+        help=("Make an alternate install, installing only the X.Y versioned "
+              "scripts (Default: pipX, pipX.Y, easy_install-X.Y)."),
+    )
+    parser.add_argument(
+        "--default-pip",
+        action="store_true",
+        default=False,
+        help=("Make a default pip install, installing the unqualified pip "
+              "and easy_install in addition to the versioned scripts."),
+    )
+
+    args = parser.parse_args(argv)
+
+    return _bootstrap(
+        root=args.root,
+        upgrade=args.upgrade,
+        user=args.user,
+        verbosity=args.verbosity,
+        altinstall=args.altinstall,
+        default_pip=args.default_pip,
+    )

Lib/ensurepip/__main__.py

@@ -0,0 +1,5 @@
+import ensurepip
+import sys
+
+if __name__ == "__main__":
+    sys.exit(ensurepip._main())

Lib/ensurepip/_uninstall.py

@@ -0,0 +1,31 @@
+"""Basic pip uninstallation support, helper for the Windows uninstaller"""
+
+import argparse
+import ensurepip
+import sys
+
+
+def _main(argv=None):
+    parser = argparse.ArgumentParser(prog="python -m ensurepip._uninstall")
+    parser.add_argument(
+        "--version",
+        action="version",
+        version="pip {}".format(ensurepip.version()),
+        help="Show the version of pip this will attempt to uninstall.",
+    )
+    parser.add_argument(
+        "-v", "--verbose",
+        action="count",
+        default=0,
+        dest="verbosity",
+        help=("Give more output. Option is additive, and can be used up to 3 "
+              "times."),
+    )
+
+    args = parser.parse_args(argv)
+
+    return ensurepip._uninstall_helper(verbosity=args.verbosity)
+
+
+if __name__ == "__main__":
+    sys.exit(_main())

Lib/test/test_ensurepip.py

@@ -0,0 +1,367 @@
+import contextlib
+import os
+import os.path
+import sys
+import tempfile
+import test.support
+import unittest
+import unittest.mock
+
+import ensurepip
+import ensurepip._uninstall
+
+
+class TestPackages(unittest.TestCase):
+    def touch(self, directory, filename):
+        fullname = os.path.join(directory, filename)
+        open(fullname, "wb").close()
+
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
+    def test_version(self):
+        # Test version()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            self.touch(tmpdir, "pip-1.2.3b1-py2.py3-none-any.whl")
+            self.touch(tmpdir, "setuptools-49.1.3-py3-none-any.whl")
+            with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None),
+                  unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir)):
+                self.assertEqual(ensurepip.version(), '1.2.3b1')
+
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
+    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()
+
+            # when bundled wheel packages are used, we get _PIP_VERSION
+            self.assertEqual(ensurepip._PIP_VERSION, ensurepip.version())
+
+        # use bundled wheel packages
+        self.assertIsNotNone(packages['pip'].wheel_name)
+        self.assertIsNotNone(packages['setuptools'].wheel_name)
+
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
+    def test_get_packages_with_dir(self):
+        # Test _get_packages() with a wheel package directory
+        setuptools_filename = "setuptools-49.1.3-py3-none-any.whl"
+        pip_filename = "pip-20.2.2-py2.py3-none-any.whl"
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            self.touch(tmpdir, setuptools_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")
+
+            with (unittest.mock.patch.object(ensurepip, '_PACKAGES', None),
+                  unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir)):
+                packages = ensurepip._get_packages()
+
+            self.assertEqual(packages['setuptools'].version, '49.1.3')
+            self.assertEqual(packages['setuptools'].wheel_path,
+                             os.path.join(tmpdir, setuptools_filename))
+            self.assertEqual(packages['pip'].version, '20.2.2')
+            self.assertEqual(packages['pip'].wheel_path,
+                             os.path.join(tmpdir, pip_filename))
+
+            # wheel package is ignored
+            self.assertEqual(sorted(packages), ['pip', 'setuptools'])
+
+
+class EnsurepipMixin:
+
+    def setUp(self):
+        run_pip_patch = unittest.mock.patch("ensurepip._run_pip")
+        self.run_pip = run_pip_patch.start()
+        self.run_pip.return_value = 0
+        self.addCleanup(run_pip_patch.stop)
+
+        # Avoid side effects on the actual os module
+        real_devnull = os.devnull
+        os_patch = unittest.mock.patch("ensurepip.os")
+        patched_os = os_patch.start()
+        # But expose os.listdir() used by _find_packages()
+        patched_os.listdir = os.listdir
+        self.addCleanup(os_patch.stop)
+        patched_os.devnull = real_devnull
+        patched_os.path = os.path
+        self.os_environ = patched_os.environ = os.environ.copy()
+
+
+class TestBootstrap(EnsurepipMixin, unittest.TestCase):
+
+    def test_basic_bootstrapping(self):
+        ensurepip.bootstrap()
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+        additional_paths = self.run_pip.call_args[0][1]
+        self.assertEqual(len(additional_paths), 2)
+
+    def test_bootstrapping_with_root(self):
+        ensurepip.bootstrap(root="/foo/bar/")
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "--root", "/foo/bar/",
+                "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+    def test_bootstrapping_with_user(self):
+        ensurepip.bootstrap(user=True)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "--user", "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+    def test_bootstrapping_with_upgrade(self):
+        ensurepip.bootstrap(upgrade=True)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "--upgrade", "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+    def test_bootstrapping_with_verbosity_1(self):
+        ensurepip.bootstrap(verbosity=1)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "-v", "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+    def test_bootstrapping_with_verbosity_2(self):
+        ensurepip.bootstrap(verbosity=2)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "-vv", "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+    def test_bootstrapping_with_verbosity_3(self):
+        ensurepip.bootstrap(verbosity=3)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "-vvv", "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+    def test_bootstrapping_with_regular_install(self):
+        ensurepip.bootstrap()
+        self.assertEqual(self.os_environ["ENSUREPIP_OPTIONS"], "install")
+
+    def test_bootstrapping_with_alt_install(self):
+        ensurepip.bootstrap(altinstall=True)
+        self.assertEqual(self.os_environ["ENSUREPIP_OPTIONS"], "altinstall")
+
+    def test_bootstrapping_with_default_pip(self):
+        ensurepip.bootstrap(default_pip=True)
+        self.assertNotIn("ENSUREPIP_OPTIONS", self.os_environ)
+
+    def test_altinstall_default_pip_conflict(self):
+        with self.assertRaises(ValueError):
+            ensurepip.bootstrap(altinstall=True, default_pip=True)
+        self.assertFalse(self.run_pip.called)
+
+    def test_pip_environment_variables_removed(self):
+        # ensurepip deliberately ignores all pip environment variables
+        # See http://bugs.python.org/issue19734 for details
+        self.os_environ["PIP_THIS_SHOULD_GO_AWAY"] = "test fodder"
+        ensurepip.bootstrap()
+        self.assertNotIn("PIP_THIS_SHOULD_GO_AWAY", self.os_environ)
+
+    def test_pip_config_file_disabled(self):
+        # ensurepip deliberately ignores the pip config file
+        # See http://bugs.python.org/issue20053 for details
+        ensurepip.bootstrap()
+        self.assertEqual(self.os_environ["PIP_CONFIG_FILE"], os.devnull)
+
+@contextlib.contextmanager
+def fake_pip(version=ensurepip.version()):
+    if version is None:
+        pip = None
+    else:
+        class FakePip():
+            __version__ = version
+        pip = FakePip()
+    sentinel = object()
+    orig_pip = sys.modules.get("pip", sentinel)
+    sys.modules["pip"] = pip
+    try:
+        yield pip
+    finally:
+        if orig_pip is sentinel:
+            del sys.modules["pip"]
+        else:
+            sys.modules["pip"] = orig_pip
+
+class TestUninstall(EnsurepipMixin, unittest.TestCase):
+
+    def test_uninstall_skipped_when_not_installed(self):
+        with fake_pip(None):
+            ensurepip._uninstall_helper()
+        self.assertFalse(self.run_pip.called)
+
+    def test_uninstall_skipped_with_warning_for_wrong_version(self):
+        with fake_pip("not a valid version"):
+            with test.support.captured_stderr() as stderr:
+                ensurepip._uninstall_helper()
+        warning = stderr.getvalue().strip()
+        self.assertIn("only uninstall a matching version", warning)
+        self.assertFalse(self.run_pip.called)
+
+
+    def test_uninstall(self):
+        with fake_pip():
+            ensurepip._uninstall_helper()
+
+        self.run_pip.assert_called_once_with(
+            [
+                "uninstall", "-y", "--disable-pip-version-check", "pip",
+                "setuptools",
+            ]
+        )
+
+    def test_uninstall_with_verbosity_1(self):
+        with fake_pip():
+            ensurepip._uninstall_helper(verbosity=1)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "uninstall", "-y", "--disable-pip-version-check", "-v", "pip",
+                "setuptools",
+            ]
+        )
+
+    def test_uninstall_with_verbosity_2(self):
+        with fake_pip():
+            ensurepip._uninstall_helper(verbosity=2)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "uninstall", "-y", "--disable-pip-version-check", "-vv", "pip",
+                "setuptools",
+            ]
+        )
+
+    def test_uninstall_with_verbosity_3(self):
+        with fake_pip():
+            ensurepip._uninstall_helper(verbosity=3)
+
+        self.run_pip.assert_called_once_with(
+            [
+                "uninstall", "-y", "--disable-pip-version-check", "-vvv",
+                "pip", "setuptools",
+            ]
+        )
+
+    def test_pip_environment_variables_removed(self):
+        # ensurepip deliberately ignores all pip environment variables
+        # See http://bugs.python.org/issue19734 for details
+        self.os_environ["PIP_THIS_SHOULD_GO_AWAY"] = "test fodder"
+        with fake_pip():
+            ensurepip._uninstall_helper()
+        self.assertNotIn("PIP_THIS_SHOULD_GO_AWAY", self.os_environ)
+
+    def test_pip_config_file_disabled(self):
+        # ensurepip deliberately ignores the pip config file
+        # See http://bugs.python.org/issue20053 for details
+        with fake_pip():
+            ensurepip._uninstall_helper()
+        self.assertEqual(self.os_environ["PIP_CONFIG_FILE"], os.devnull)
+
+
+# Basic testing of the main functions and their argument parsing
+
+EXPECTED_VERSION_OUTPUT = "pip " + ensurepip.version()
+
+class TestBootstrappingMainFunction(EnsurepipMixin, unittest.TestCase):
+
+    def test_bootstrap_version(self):
+        with test.support.captured_stdout() as stdout:
+            with self.assertRaises(SystemExit):
+                ensurepip._main(["--version"])
+        result = stdout.getvalue().strip()
+        self.assertEqual(result, EXPECTED_VERSION_OUTPUT)
+        self.assertFalse(self.run_pip.called)
+
+    def test_basic_bootstrapping(self):
+        exit_code = ensurepip._main([])
+
+        self.run_pip.assert_called_once_with(
+            [
+                "install", "--no-cache-dir", "--no-index", "--find-links",
+                unittest.mock.ANY, "setuptools", "pip",
+            ],
+            unittest.mock.ANY,
+        )
+
+        additional_paths = self.run_pip.call_args[0][1]
+        self.assertEqual(len(additional_paths), 2)
+        self.assertEqual(exit_code, 0)
+
+    def test_bootstrapping_error_code(self):
+        self.run_pip.return_value = 2
+        exit_code = ensurepip._main([])
+        self.assertEqual(exit_code, 2)
+
+
+class TestUninstallationMainFunction(EnsurepipMixin, unittest.TestCase):
+
+    def test_uninstall_version(self):
+        with test.support.captured_stdout() as stdout:
+            with self.assertRaises(SystemExit):
+                ensurepip._uninstall._main(["--version"])
+        result = stdout.getvalue().strip()
+        self.assertEqual(result, EXPECTED_VERSION_OUTPUT)
+        self.assertFalse(self.run_pip.called)
+
+    def test_basic_uninstall(self):
+        with fake_pip():
+            exit_code = ensurepip._uninstall._main([])
+
+        self.run_pip.assert_called_once_with(
+            [
+                "uninstall", "-y", "--disable-pip-version-check", "pip",
+                "setuptools",
+            ]
+        )
+
+        self.assertEqual(exit_code, 0)
+
+    def test_uninstall_error_code(self):
+        with fake_pip():
+            self.run_pip.return_value = 2
+            exit_code = ensurepip._uninstall._main([])
+        self.assertEqual(exit_code, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()