Improve handling of non-normalized .dist-info folders (#168)

Co-authored-by: Pradyun Gedam <[email protected]>

Author: pradyunsg

src/installer/sources.py

@@ -5,8 +5,9 @@
 import stat
 import zipfile
 from contextlib import contextmanager
-from typing import BinaryIO, ClassVar, Iterator, List, Tuple, Type, cast
+from typing import BinaryIO, ClassVar, Iterator, List, Optional, Tuple, Type, cast
 
+from installer.exceptions import InstallerError
 from installer.records import RecordEntry, parse_record_file
 from installer.utils import canonicalize_name, parse_wheel_filename
 
@@ -101,17 +102,32 @@ def get_contents(self) -> Iterator[WheelContentElement]:
         raise NotImplementedError
 
 
-class _WheelFileValidationError(ValueError):
+class _WheelFileValidationError(ValueError, InstallerError):
     """Raised when a wheel file fails validation."""
 
-    def __init__(self, issues: List[str]) -> None:  # noqa: D107
+    def __init__(self, issues: List[str]) -> None:
         super().__init__(repr(issues))
         self.issues = issues
 
     def __repr__(self) -> str:
         return f"WheelFileValidationError(issues={self.issues!r})"
 
 
+class _WheelFileBadDistInfo(ValueError, InstallerError):
+    """Raised when a wheel file has issues around `.dist-info`."""
+
+    def __init__(self, *, reason: str, filename: Optional[str], dist_info: str) -> None:
+        super().__init__(reason)
+        self.reason = reason
+        self.filename = filename
+        self.dist_info = dist_info
+
+    def __str__(self) -> str:
+        return (
+            f"{self.reason} (filename={self.filename!r}, dist_info={self.dist_info!r})"
+        )
+
+
 class WheelFile(WheelSource):
     """Implements `WheelSource`, for an existing file from the filesystem.
 
@@ -137,6 +153,7 @@ def __init__(self, f: zipfile.ZipFile) -> None:
             version=parsed_name.version,
             distribution=parsed_name.distribution,
         )
+        self._dist_info_dir: Optional[str] = None
 
     @classmethod
     @contextmanager
@@ -148,29 +165,39 @@ def open(cls, path: "os.PathLike[str]") -> Iterator["WheelFile"]:
     @property
     def dist_info_dir(self) -> str:
         """Name of the dist-info directory."""
-        if not hasattr(self, "_dist_info_dir"):
-            top_level_directories = {
-                path.split("/", 1)[0] for path in self._zipfile.namelist()
-            }
-            dist_infos = [
-                name for name in top_level_directories if name.endswith(".dist-info")
-            ]
-
-            assert (
-                len(dist_infos) == 1
-            ), "Wheel doesn't contain exactly one .dist-info directory"
-            dist_info_dir = dist_infos[0]
-
-            # NAME-VER.dist-info
-            di_dname = dist_info_dir.rsplit("-", 2)[0]
-            norm_di_dname = canonicalize_name(di_dname)
-            norm_file_dname = canonicalize_name(self.distribution)
-            assert (
-                norm_di_dname == norm_file_dname
-            ), "Wheel .dist-info directory doesn't match wheel filename"
-
-            self._dist_info_dir = dist_info_dir
-        return self._dist_info_dir
+        if self._dist_info_dir is not None:
+            return self._dist_info_dir
+
+        top_level_directories = {
+            path.split("/", 1)[0] for path in self._zipfile.namelist()
+        }
+        dist_infos = [
+            name for name in top_level_directories if name.endswith(".dist-info")
+        ]
+
+        try:
+            (dist_info_dir,) = dist_infos
+        except ValueError:
+            raise _WheelFileBadDistInfo(
+                reason="Wheel doesn't contain exactly one .dist-info directory",
+                filename=self._zipfile.filename,
+                dist_info=str(sorted(dist_infos)),
+            ) from None
+
+        # NAME-VER.dist-info
+        di_dname = dist_info_dir.rsplit("-", 2)[0]
+        norm_di_dname = canonicalize_name(di_dname)
+        norm_file_dname = canonicalize_name(self.distribution)
+
+        if norm_di_dname != norm_file_dname:
+            raise _WheelFileBadDistInfo(
+                reason="Wheel .dist-info directory doesn't match wheel filename",
+                filename=self._zipfile.filename,
+                dist_info=dist_info_dir,
+            )
+
+        self._dist_info_dir = dist_info_dir
+        return dist_info_dir
 
     @property
     def dist_info_filenames(self) -> List[str]:

tests/test_sources.py

@@ -6,6 +6,7 @@
 
 import pytest
 
+from installer.exceptions import InstallerError
 from installer.records import parse_record_file
 from installer.sources import WheelFile, WheelSource
 
@@ -133,10 +134,35 @@ def test_requires_dist_info_name_match(self, fancy_wheel):
         )
         # Python 3.7: rename doesn't return the new name:
         misnamed = fancy_wheel.parent / "misnamed-1.0.0-py3-none-any.whl"
-        with pytest.raises(AssertionError):
+        with pytest.raises(InstallerError) as ctx:
             with WheelFile.open(misnamed) as source:
                 source.dist_info_filenames
 
+        error = ctx.value
+        print(error)
+        assert error.filename == str(misnamed)
+        assert error.dist_info == "fancy-1.0.0.dist-info"
+        assert "" in error.reason
+        assert error.dist_info in str(error)
+
+    def test_enforces_single_dist_info(self, fancy_wheel):
+        with zipfile.ZipFile(fancy_wheel, "a") as archive:
+            archive.writestr(
+                "name-1.0.0.dist-info/random.txt",
+                b"This is a random file.",
+            )
+
+        with pytest.raises(InstallerError) as ctx:
+            with WheelFile.open(fancy_wheel) as source:
+                source.dist_info_filenames
+
+        error = ctx.value
+        print(error)
+        assert error.filename == str(fancy_wheel)
+        assert error.dist_info == str(["fancy-1.0.0.dist-info", "name-1.0.0.dist-info"])
+        assert "exactly one .dist-info" in error.reason
+        assert error.dist_info in str(error)
+
     def test_rejects_no_record_on_validate(self, fancy_wheel):
         # Remove RECORD
         replace_file_in_zip(