8000 Use `pathlib` in `ensurepip` internally · python/cpython@f3a49c2 · GitHub
[go: up one dir, main page]
Skip to content

Commit f3a49c2

Browse files
webknjazwebknjaz
committed
Use pathlib in ensurepip internally
It provides us with the ability to write simpler high-level logic that is easier to understand. As a side effect, it became possible to make the code less branchy.
1 parent 7b04736 commit f3a49c2

File tree

2 files changed

+60
-58
lines changed

2 files changed

+60
-58
lines changed

‎Lib/ensurepip/__init__.py

Lines changed: 48 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
import os
2-
import os.path
32
import subprocess
43
import sys
54
import sysconfig
65
import tempfile
6+
from contextlib import nullcontext, suppress
77
from importlib import resources
8+
from pathlib import Path
9+
from shutil import copy2
810

911

1012
__all__ = ["version", "bootstrap"]
@@ -14,39 +16,51 @@
1416
# policies recommend against bundling dependencies. For example, Fedora
1517
# installs wheel packages in the /usr/share/python-wheels/ directory and don't
1618
# install the ensurepip._bundled package.
17-
_WHEEL_PKG_DIR = sysconfig.get_config_var('WHEEL_PKG_DIR')
19+
_WHEEL_PKG_DIR = (
20+
whl_pkg_dir_str := sysconfig.get_config_var('WHEEL_PKG_DIR')
21+
) is not None and Path(whl_pkg_dir_str).resolve() or None
1822

1923

2024
def _find_wheel_pkg_dir_pip():
2125
if _WHEEL_PKG_DIR is None:
22-
return None
26+
raise LookupError(
27+
'The compile-time `WHEEL_PKG_DIR` is unset so there is '
28+
'no place for looking up the wheels.',
29+
) from None
30+
31+
32+
dist_matching_wheels = _WHEEL_PKG_DIR.glob(f'pip-*.whl')
2333
try:
24-
filenames = os.listdir(_WHEEL_PKG_DIR)
25-
except OSError:
26-
# Ignore: path doesn't exist or permission error
27-
return None
28-
# Make the code deterministic if a directory contains multiple wheel files
29-
# of the same package, but don't attempt to implement correct version
30-
# comparison since this case should not happen.
31-
filenames = sorted(filenames, reverse=True)
32-
for filename in filenames:
33-
# filename is like 'pip-21.2.4-py3-none-any.whl'
34-
if not filename.startswith("pip-") or not filename.endswith(".whl"):
35-
continue
36-
37-
# Extract '21.2.4' from 'pip-21.2.4-py3-none-any.whl'
38-
version = filename.removeprefix("pip-").partition("-")[0]
39-
return {"version": version, "filename": filename, "bundled": False}
40-
41-
return None
42-
43-
44-
def _get_pip_info():
34+
last_matching_dist_wheel = sorted(dist_matching_wheels)[-1]
35+
except IndexError as index_err:
36+
raise LookupError(
37+
'`WHEEL_PKG_DIR` does not contain any wheel files for `pip`.',
38+
) from index_err
39+
40+
return nullcontext(last_matching_dist_wheel)
41+
42+
43+
def _get_pip_whl_path_ctx():
4544
# Prefer pip from the wheel package directory, if present.
46-
if (pip_info := _find_wheel_pkg_dir_pip()) is not None:
47-
return pip_info
48-
filename = f"pip-{_PIP_VERSION}-py3-none-any.whl"
49-
return {"version": _PIP_VERSION, "filename": filename, "bundled": True}
45+
with suppress(LookupError):
46+
return _find_wheel_pkg_dir_pip()
47+
48+
return resources.as_file(
49+
resources.files('ensurepip')
50+
/ '_bundled'
51+
/ f'pip-{_PIP_VERSION}-py3-none-any.whl'
52+
)
53+
54+
55+
def _get_pip_version():
56+
with _get_pip_whl_path_ctx() as bundled_wheel_path:
57+
wheel_name = bundled_wheel_path.name
58+
return (
59+
# Extract '21.2.4' from 'pip-21.2.4-py3-none-any.whl'
60+
wheel_name.
61+
removeprefix('pip-').
62+
partition('-')[0]
63+
)
5064

5165

5266
def _run_pip(args, additional_paths=None):
@@ -79,7 +93,7 @@ def version():
7993
"""
8094
Returns a string specifying the bundled version of pip.
8195
"""
82-
return _get_pip_info()["version"]
96+
return _get_pip_version()
8397

8498

8599
def _disable_pip_configuration_settings():
@@ -141,20 +155,10 @@ def _bootstrap(*, root=None, upgrade=False, user=False,
141155
with tempfile.TemporaryDirectory() as tmpdir:
142156
# Put our bundled wheels into a temporary directory and construct the
143157
# additional paths that need added to sys.path
144-
package = _get_pip_info()
145-
wheel_name = package["filename"]
146-
if package["bundled"]:
147-
# Use bundled wheel package
148-
wheel_path = resources.files("ensurepip") / "_bundled" / wheel_name
149-
whl = wheel_path.read_bytes()
150-
else:
151-
# Use the wheel package directory
152-
with open(os.path.join(_WHEEL_PKG_DIR, wheel_name), "rb") as fp:
153-
whl = fp.read()
154-
155-
filename = os.path.join(tmpdir, wheel_name)
156-
with open(filename, "wb") as fp:
157-
fp.write(whl)
158+
tmpdir_path = Path(tmpdir)
159+
with _get_pip_whl_path_ctx() as bundled_wheel_path:
160+
tmp_wheel_path = tmpdir_path / bundled_wheel_path.name
161+
copy2(bundled_wheel_path, tmp_wheel_path)
158162

159163
# Construct the arguments to be passed to the pip command
160164
args = ["install", "--no-cache-dir", "--no-index", "--find-links", tmpdir]
@@ -167,7 +171,7 @@ def _bootstrap(*, root=None, upgrade=False, user=False,
167171
if verbosity:
168172
args += ["-" + "v" * verbosity]
169173

170-
return _run_pip([*args, "pip"], [filename])
174+
return _run_pip([*args, "pip"], [os.fsencode(tmp_wheel_path)])
171175

172176

173177
def _uninstall_helper(*, verbosity=0):

‎Lib/test/test_ensurepip.py

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
import test.support
77
import unittest
88
import unittest.mock
9+
from importlib.resources.abc import Traversable
10+
from pathlib import Path
911

1012
import ensurepip
1113
import ensurepip._uninstall
@@ -20,22 +22,20 @@ def test_version(self):
2022
# Test version()
2123
with tempfile.TemporaryDirectory() as tmpdir:
2224
self.touch(tmpdir, "pip-1.2.3b1-py2.py3-none-any.whl")
23-
with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir):
25+
with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', Path(tmpdir)):
2426
self.assertEqual(ensurepip.version(), '1.2.3b1')
2527

26-
def test_get_pip_info_no_dir(self):
27-
# Test _get_pip_info() without a wheel package directory
28+
def test_version_no_dir(self):
29+
# Test version() without a wheel package directory
2830
with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', None):
29-
pip_info = ensurepip._get_pip_info()
30-
3131
# when the bundled pip wheel is used, we get _PIP_VERSION
3232
self.assertEqual(ensurepip._PIP_VERSION, ensurepip.version())
3333

34-
# use the bundled pip wheel
34+
def test_selected_wheel_path_no_dir(self):
3535
pip_filename = f'pip-{ensurepip._PIP_VERSION}-py3-none-any.whl'
36-
expected = {"version": ensurepip._PIP_VERSION, "filename": pip_filename,
37-
"bundled": True}
38-
self.assertDictEqual(pip_info, expected)
36+
with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', None):
37+
with ensurepip._get_pip_whl_path_ctx() as bundled_wheel_path:
38+
self.assertEqual(pip_filename, bundled_wheel_path.name)
3939

4040
def test_get_pip_info_with_dir(self):
4141
# Test _get_pip_info() with a wheel package directory
@@ -48,11 +48,9 @@ def test_get_pip_info_with_dir(self):
4848
self.touch(tmpdir, "wheel-0.34.2-py2.py3-none-any.whl")
4949
self.touch(tmpdir, "pip-script.py")
5050

51-
with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', tmpdir):
52-
pip_info = ensurepip._get_pip_info()
53-
54-
expected = {"version": '20.2.2', "filename": pip_filename, "bundled": False}
55-
self.assertDictEqual(pip_info, expected)
51+
with unittest.mock.patch.object(ensurepip, '_WHEEL_PKG_DIR', Path(tmpdir)):
52+
with ensurepip._get_pip_whl_path_ctx() as bundled_wheel_path:
53+
self.assertEqual(pip_filename, bundled_wheel_path.name)
5654

5755

5856
class EnsurepipMixin:

0 commit comments

Comments
 (0)
0