Refactor wheel info extraction during install

src/pip/_internal/operations/install/wheel.py

@@ -23,7 +23,7 @@
 from pip._vendor.distlib.scripts import ScriptMaker
 from pip._vendor.distlib.util import get_export_entry
 from pip._vendor.packaging.utils import canonicalize_name
-from pip._vendor.six import StringIO
+from pip._vendor.six import StringIO, ensure_str
 
 from pip._internal.exceptions import InstallationError, UnsupportedWheel
 from pip._internal.locations import get_major_minor_version
@@ -33,6 +33,7 @@
 from pip._internal.utils.unpacking import unpack_file
 
 if MYPY_CHECK_RUNNING:
+    from email.message import Message
     from typing import (
         Dict, List, Optional, Sequence, Tuple, IO, Text, Any,
         Iterable, Callable, Set,
@@ -97,23 +98,9 @@ def fix_script(path):
     return None
 
 
-dist_info_re = re.compile(r"""^(?P<namever>(?P<name>.+?)(-(?P<ver>.+?))?)
-                                \.dist-info$""", re.VERBOSE)
-
-
-def root_is_purelib(name, wheeldir):
-    # type: (str, str) -> bool
-    """True if the extracted wheel in wheeldir should go into purelib."""
-    name_folded = name.replace("-", "_")
-    for item in os.listdir(wheeldir):
-        match = dist_info_re.match(item)
-        if match and match.group('name') == name_folded:
-            with open(os.path.join(wheeldir, item, 'WHEEL')) as wheel:
-                for line in wheel:
-                    line = line.lower().rstrip()
-                    if line == "root-is-purelib: true":
-                        return True
-    return False
+def wheel_root_is_purelib(metadata):
+    # type: (Message) -> bool
+    return metadata.get("Root-Is-Purelib", "").lower() == "true"
 
 
 def get_entrypoints(filename):
@@ -324,23 +311,25 @@ def install_unpacked_wheel(
     # TODO: Look into moving this into a dedicated class for representing an
     #       installation.
 
+    source = wheeldir.rstrip(os.path.sep) + os.path.sep
+
     try:
-        version = wheel_version(wheeldir)
+        info_dir = wheel_dist_info_dir(source, name)
+        metadata = wheel_metadata(source, info_dir)
+        version = wheel_version(metadata)
     except UnsupportedWheel as e:
         raise UnsupportedWheel(
             "{} has an invalid wheel, {}".format(name, str(e))
         )
 
     check_compatibility(version, name)
 
-    if root_is_purelib(name, wheeldir):
+    if wheel_root_is_purelib(metadata):
         lib_dir = scheme.purelib
     else:
         lib_dir = scheme.platlib
 
-    source = wheeldir.rstrip(os.path.sep) + os.path.sep
     subdirs = os.listdir(source)
-    info_dirs = [s for s in subdirs if s.endswith('.dist-info')]
     data_dirs = [s for s in subdirs if s.endswith('.data')]
 
     # Record details of the files moved
@@ -434,27 +423,6 @@ def clobber(
 
     clobber(source, lib_dir, True)
 
-    assert info_dirs, "{} .dist-info directory not found".format(
-        req_description
-    )
-
-    assert len(info_dirs) == 1, (
-        '{} multiple .dist-info directories found: {}'.format(
-            req_description, ', '.join(info_dirs)
-        )
-    )
-
-    info_dir = info_dirs[0]
-
-    info_dir_name = canonicalize_name(info_dir)
-    canonical_name = canonicalize_name(name)
-    if not info_dir_name.startswith(canonical_name):
-        raise UnsupportedWheel(
-            "{} .dist-info directory {!r} does not start with {!r}".format(
-                req_description, info_dir, canonical_name
-            )
-        )
-
     dest_info_dir = os.path.join(lib_dir, info_dir)
 
     # Get the defined entry points
@@ -656,25 +624,48 @@ def install_wheel(
         )
 
 
-def wheel_version(source_dir):
-    # type: (Optional[str]) -> Tuple[int, ...]
-    """Return the Wheel-Version of an extracted wheel, if possible.
-    Otherwise, raise UnsupportedWheel if we couldn't parse / extract it.
+def wheel_dist_info_dir(source, name):
+    # type: (str, str) -> str
+    """Returns the name of the contained .dist-info directory.
+
+    Raises AssertionError or UnsupportedWheel if not found, >1 found, or
+    it doesn't match the provided name.
     """
-    try:
-        dists = [d for d in pkg_resources.find_on_path(None, source_dir)]
-    except Exception as e:
+    subdirs = os.listdir(source)
+    info_dirs = [s for s in subdirs if s.endswith('.dist-info')]
+
+    if not info_dirs:
+        raise UnsupportedWheel(".dist-info directory not found")
+
+    if len(info_dirs) > 1:
+        raise UnsupportedWheel(
+            "multiple .dist-info directories found: {}".format(
+                ", ".join(info_dirs)
+            )
+        )
+
+    info_dir = info_dirs[0]
+
+    info_dir_name = canonicalize_name(info_dir)
+    canonical_name = canonicalize_name(name)
+    if not info_dir_name.startswith(canonical_name):
         raise UnsupportedWheel(
-            "could not find a contained distribution due to: {!r}".format(e)
+            ".dist-info directory {!r} does not start with {!r}".format(
+                info_dir, canonical_name
+            )
         )
 
-    if not dists:
-        raise UnsupportedWheel("no contained distribution found")
+    return info_dir
 
-    dist = dists[0]
 
+def wheel_metadata(source, dist_info_dir):
+    # type: (str, str) -> Message
+    """Return the WHEEL metadata of an extracted wheel, if possible.
+    Otherwise, raise UnsupportedWheel.
+    """
     try:
-        wheel_text = dist.get_metadata('WHEEL')
+        with open(os.path.join(source, dist_info_dir, "WHEEL"), "rb") as f:
+            wheel_text = ensure_str(f.read())
     except (IOError, OSError) as e:
         raise UnsupportedWheel("could not read WHEEL file: {!r}".format(e))
     except UnicodeDecodeError as e:
@@ -683,8 +674,14 @@ def wheel_version(source_dir):
     # FeedParser (used by Parser) does not raise any exceptions. The returned
     # message may have .defects populated, but for backwards-compatibility we
     # currently ignore them.
-    wheel_data = Parser().parsestr(wheel_text)
+    return Parser().parsestr(wheel_text)
 
+
+def wheel_version(wheel_data):
+    # type: (Message) -> Tuple[int, ...]
+    """Given WHEEL metadata, return the parsed Wheel-Version.
+    Otherwise, raise UnsupportedWheel.
+    """
     version_text = wheel_data["Wheel-Version"]
     if version_text is None:
         raise UnsupportedWheel("WHEEL is missing Wheel-Version")

tests/unit/test_wheel.py

@@ -3,10 +3,10 @@
 import logging
 import os
 import textwrap
+from email import message_from_string
 
 import pytest
-from mock import Mock, patch
-from pip._vendor import pkg_resources
+from mock import patch
 from pip._vendor.packaging.requirements import Requirement
 
 from pip._internal.exceptions import UnsupportedWheel
@@ -190,59 +190,64 @@ def test_get_csv_rows_for_installed__long_lines(tmpdir, caplog):
     assert messages == expected
 
 
-def test_wheel_version(tmpdir, data):
-    future_wheel = 'futurewheel-1.9-py2.py3-none-any.whl'
-    future_version = (1, 9)
+def test_wheel_dist_info_dir_found(tmpdir):
+    expected = "simple-0.1.dist-info"
+    tmpdir.joinpath(expected).mkdir()
+    assert wheel.wheel_dist_info_dir(str(tmpdir), "simple") == expected
 
-    unpack_file(data.packages.joinpath(future_wheel), tmpdir + 'future')
 
-    assert wheel.wheel_version(tmpdir + 'future') == future_version
+def test_wheel_dist_info_dir_multiple(tmpdir):
+    tmpdir.joinpath("simple-0.1.dist-info").mkdir()
+    tmpdir.joinpath("unrelated-0.1.dist-info").mkdir()
+    with pytest.raises(UnsupportedWheel) as e:
+        wheel.wheel_dist_info_dir(str(tmpdir), "simple")
+    assert "multiple .dist-info directories found" in str(e.value)
 
 
-def test_wheel_version_fails_on_error(monkeypatch):
-    err = RuntimeError("example")
-    monkeypatch.setattr(pkg_resources, "find_on_path", Mock(side_effect=err))
+def test_wheel_dist_info_dir_none(tmpdir):
     with pytest.raises(UnsupportedWheel) as e:
-        wheel.wheel_version(".")
-    assert repr(err) in str(e.value)
+        wheel.wheel_dist_info_dir(str(tmpdir), "simple")
+    assert "directory not found" in str(e.value)
 
 
-def test_wheel_version_fails_no_dists(tmpdir):
+def test_wheel_dist_info_dir_wrong_name(tmpdir):
+    tmpdir.joinpath("unrelated-0.1.dist-info").mkdir()
     with pytest.raises(UnsupportedWheel) as e:
-        wheel.wheel_version(str(tmpdir))
-    assert "no contained distribution found" in str(e.value)
+        wheel.wheel_dist_info_dir(str(tmpdir), "simple")
+    assert "does not start with 'simple'" in str(e.value)
+
+
+def test_wheel_version_ok(tmpdir, data):
+    assert wheel.wheel_version(
+        message_from_string("Wheel-Version: 1.9")
+    ) == (1, 9)
 
 
-def test_wheel_version_fails_missing_wheel(tmpdir):
+def test_wheel_metadata_fails_missing_wheel(tmpdir):
     dist_info_dir = tmpdir / "simple-0.1.0.dist-info"
     dist_info_dir.mkdir()
     dist_info_dir.joinpath("METADATA").touch()
 
     with pytest.raises(UnsupportedWheel) as e:
-        wheel.wheel_version(str(tmpdir))
+        wheel.wheel_metadata(str(tmpdir), dist_info_dir.name)
     assert "could not read WHEEL file" in str(e.value)
 
 
 @skip_if_python2
-def test_wheel_version_fails_on_bad_encoding(tmpdir):
+def test_wheel_metadata_fails_on_bad_encoding(tmpdir):
     dist_info_dir = tmpdir / "simple-0.1.0.dist-info"
     dist_info_dir.mkdir()
     dist_info_dir.joinpath("METADATA").touch()
     dist_info_dir.joinpath("WHEEL").write_bytes(b"\xff")
 
     with pytest.raises(UnsupportedWheel) as e:
-        wheel.wheel_version(str(tmpdir))
+        wheel.wheel_metadata(str(tmpdir), dist_info_dir.name)
     assert "error decoding WHEEL" in str(e.value)
 
 
-def test_wheel_version_fails_on_no_wheel_version(tmpdir):
-    dist_info_dir = tmpdir / "simple-0.1.0.dist-info"
-    dist_info_dir.mkdir()
-    dist_info_dir.joinpath("METADATA").touch()
-    dist_info_dir.joinpath("WHEEL").touch()
-
+def test_wheel_version_fails_on_no_wheel_version():
     with pytest.raises(UnsupportedWheel) as e:
-        wheel.wheel_version(str(tmpdir))
+        wheel.wheel_version(message_from_string(""))
     assert "missing Wheel-Version" in str(e.value)
 
 
@@ -251,19 +256,26 @@ def test_wheel_version_fails_on_no_wheel_version(tmpdir):
     ("1.b",),
     ("1.",),
 ])
-def test_wheel_version_fails_on_bad_wheel_version(tmpdir, version):
-    dist_info_dir = tmpdir / "simple-0.1.0.dist-info"
-    dist_info_dir.mkdir()
-    dist_info_dir.joinpath("METADATA").touch()
-    dist_info_dir.joinpath("WHEEL").write_text(
-        "Wheel-Version: {}".format(version)
-    )
-
+def test_wheel_version_fails_on_bad_wheel_version(version):
     with pytest.raises(UnsupportedWheel) as e:
-        wheel.wheel_version(str(tmpdir))
+        wheel.wheel_version(
+            message_from_string("Wheel-Version: {}".format(version))
+        )
     assert "invalid Wheel-Version" in str(e.value)
 
 
+@pytest.mark.parametrize("text,expected", [
+    ("Root-Is-Purelib: true", True),
+    ("Root-Is-Purelib: false", False),
+    ("Root-Is-Purelib: hello", False),
+    ("", False),
+    ("root-is-purelib: true", True),
+    ("root-is-purelib: True", True),
+])
+def test_wheel_root_is_purelib(text, expected):
+    assert wheel.wheel_root_is_purelib(message_from_string(text)) == expected
+
+
 def test_check_compatibility():
     name = 'test'
     vc = wheel.VERSION_COMPATIBLE
@@ -296,22 +308,6 @@ def test_unpack_wheel_no_flatten(self, tmpdir):
         unpack_file(filepath, tmpdir)
         assert os.path.isdir(os.path.join(tmpdir, 'meta-1.0.dist-info'))
 
-    def test_purelib_platlib(self, data):
-        """
-        Test the "wheel is purelib/platlib" code.
-        """
-        packages = [
-            ("pure_wheel", data.packages.joinpath("pure_wheel-1.7"), True),
-            ("plat_wheel", data.packages.joinpath("plat_wheel-1.7"), False),
-            ("pure_wheel", data.packages.joinpath(
-                "pure_wheel-_invalidversion_"), True),
-            ("plat_wheel", data.packages.joinpath(
-                "plat_wheel-_invalidversion_"), False),
-        ]
-
-        for name, path, expected in packages:
-            assert wheel.root_is_purelib(name, path) == expected
-
 
 class TestInstallUnpackedWheel(object):
     """