Added compatibility shim for wheel.wheelfile

2 files changed, 124 insertions, 4 deletions

diff --git a/src/wheel/_wheelfile.py b/src/wheel/_wheelfile.py
index e1c263f..4a3190b 100644
--- a/src/wheel/_wheelfile.py
+++ b/src/wheel/_wheelfile.py
@@ -193,7 +193,11 @@ class WheelReader:
 
     def _read_record(self) -> OrderedDict[str, WheelRecordEntry]:
         entries = OrderedDict()
-        contents = self.read_dist_info("RECORD")
+        try:
+            contents = self.read_dist_info("RECORD")
+        except WheelError:
+            raise WheelError(f"Missing {self._dist_info_dir}/RECORD file") from None
+
         reader = csv.reader(
             contents.strip().split("\n"),
             delimiter=",",
@@ -310,17 +314,17 @@ class WheelReader:
             self._zip.open(archive_name), archive_name, record_entry
         )
 
-    def _read_file(self, archive_name: str) -> bytes:
+    def read_file(self, archive_name: str) -> bytes:
         with self._open_file(archive_name) as fp:
             return fp.read()
 
     def read_data_file(self, filename: str) -> bytes:
         archive_path = self._data_dir + "/" + filename.strip("/")
-        return self._read_file(archive_path)
+        return self.read_file(archive_path)
 
     def read_distinfo_file(self, filename: str) -> bytes:
         archive_path = self._dist_info_dir + "/" + filename.strip("/")
-        return self._read_file(archive_path)
+        return self.read_file(archive_path)
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.path_or_fd})"
diff --git a/src/wheel/wheelfile.py b/src/wheel/wheelfile.py
new file mode 100644
index 0000000..22c4d84
--- /dev/null
+++ b/src/wheel/wheelfile.py
@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import os.path
+import re
+import time
+from os import PathLike
+from types import TracebackType
+from typing import TYPE_CHECKING
+from warnings import warn
+from zipfile import ZipInfo
+
+from . import WheelWriter
+from ._wheelfile import DEFAULT_TIMESTAMP, WheelReader
+
+if TYPE_CHECKING:
+    from typing import Literal
+
+warn(
+    DeprecationWarning(
+        f"The {__name__} module has been deprecated in favor of a supported public "
+        "API, and will be removed in a future release."
+    )
+)
+
+WHEEL_INFO_RE = re.compile(
+    r"""^(?P<namever>(?P<name>[^\s-]+?)-(?P<ver>[^\s-]+?))(-(?P<build>\d[^\s-]*))?
+     -(?P<pyver>[^\s-]+?)-(?P<abi>[^\s-]+?)-(?P<plat>\S+)\.whl$""",
+    re.VERBOSE,
+)
+MINIMUM_TIMESTAMP = 315532800  # 1980-01-01 00:00:00 UTC
+
+
+def get_zipinfo_datetime(timestamp=None):
+    # Some applications need reproducible .whl files, but they can't do this without
+    # forcing the timestamp of the individual ZipInfo objects. See issue #143.
+    timestamp = int(os.environ.get("SOURCE_DATE_EPOCH", timestamp or time.time()))
+    timestamp = max(timestamp, MINIMUM_TIMESTAMP)
+    return time.gmtime(timestamp)[0:6]
+
+
+class WheelFile:
+    """Compatibility shim for WheelReader and WheelWriter."""
+
+    _reader: WheelReader
+    _writer: WheelWriter
+
+    def __init__(self, path: str | PathLike[str], mode: Literal["r", "w"] = "r"):
+        if mode == "r":
+            self._reader = WheelReader(path)
+        elif mode == "w":
+            self._writer = WheelWriter(path)
+        else:
+            raise ValueError(f"Invalid mode: {mode}")
+
+        self.filename = str(path)
+        self.parsed_filename = WHEEL_INFO_RE.match(os.path.basename(self.filename))
+        self.dist_info_path = f"{self.parsed_filename.group('namever')}.dist-info"
+
+    def __enter__(self) -> WheelFile:
+        if hasattr(self, "_reader"):
+            self._reader.__enter__()
+        else:
+            self._writer.__enter__()
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException],
+        exc_val: BaseException,
+        exc_tb: TracebackType,
+    ) -> None:
+        if hasattr(self, "_reader"):
+            self._reader.__exit__(exc_type, exc_val, exc_tb)
+        else:
+            self._writer.__exit__(exc_type, exc_val, exc_tb)
+
+    def read(self, name: str) -> bytes:
+        return self._reader.read_file(name)
+
+    def extractall(self, base_path: str | PathLike[str] | None = None) -> None:
+        self._reader.extractall(base_path)
+
+    def write_files(self, base_dir: PathLike[str] | str) -> None:
+        self._writer.write_files_from_directory(base_dir)
+
+    def write(
+        self,
+        filename: str | PathLike[str],
+        arcname: str | None = None,
+        compress_type: int | None = None,
+    ):
+        arcname = arcname or filename
+        self._writer.write_file(arcname, filename)
+
+    def writestr(
+        self,
+        zinfo_or_arcname: str | ZipInfo,
+        data: bytes | str,
+        compress_type: int | None = None,
+    ):
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+
+        if isinstance(zinfo_or_arcname, ZipInfo):
+            arcname = zinfo_or_arcname.filename
+            timestamp = zinfo_or_arcname.date_time
+        elif isinstance(zinfo_or_arcname, str):
+            arcname = zinfo_or_arcname
+            timestamp = DEFAULT_TIMESTAMP
+        else:
+            raise TypeError(
+                f"Invalid type for zinfo_or_arcname: {type(zinfo_or_arcname)}"
+            )
+
+        self._writer.write_file(arcname, data, timestamp=timestamp)