add tests

Author: lalten

python/private/pypi/BUILD.bazel

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load("//python:py_library.bzl", "py_library")
 
 package(default_visibility = ["//:__subpackages__"])
 
@@ -377,6 +378,12 @@ bzl_library(
     ],
 )
 
+py_library(
+    name = "repack_whl",
+    srcs = ["repack_whl.py"],
+    deps = ["//tools:wheelmaker"],
+)
+
 bzl_library(
     name = "requirements_files_by_platform_bzl",
     srcs = ["requirements_files_by_platform.bzl"],

python/private/pypi/repack_whl.py

@@ -44,6 +44,16 @@
 _DISTINFO = "dist-info"
 
 
+def _has_all_quoted_filenames(record_contents: str) -> bool:
+    """Check if all filenames in the RECORD are quoted.
+
+    Some wheels (like torch) have all filenames quoted in their RECORD file.
+    We detect this to preserve the quoting style when repacking.
+    """
+    lines = record_contents.splitlines()
+    return all(line.startswith('"') for line in lines)
+
+
 def _unidiff_output(expected, actual, record):
     """
     Helper function. Returns a string containing the unified diff of two
@@ -151,7 +161,7 @@ def main(sys_argv):
         logging.debug(f"Found dist-info dir: {distinfo_dir}")
         record_path = distinfo_dir / "RECORD"
         record_contents = record_path.read_text() if record_path.exists() else ""
-        quote_files = all(line.startswith('"') for line in record_contents.splitlines())
+        quote_files = _has_all_quoted_filenames(record_contents)
         distribution_prefix = distinfo_dir.with_suffix("").name
 
         with _WhlFile(

tests/pypi/repack_whl/BUILD.bazel

@@ -0,0 +1,8 @@
+load("//python:py_test.bzl", "py_test")
+
+py_test(
+    name = "repack_whl_test",
+    size = "small",
+    srcs = ["repack_whl_test.py"],
+    deps = ["//python/private/pypi:repack_whl"],
+)

tests/pypi/repack_whl/repack_whl_test.py

@@ -0,0 +1,37 @@
+import unittest
+
+from python.private.pypi import repack_whl
+
+
+class HasAllQuotedFilenamesTest(unittest.TestCase):
+    """Tests for _has_all_quoted_filenames detection logic."""
+
+    def test_all_quoted(self) -> None:
+        """Returns True when all lines start with quotes (torch-style)."""
+        record = """\
+"torch/__init__.py",sha256=abc,123
+"torch/utils.py",sha256=def,456
+"torch-2.0.0.dist-info/WHEEL",sha256=ghi,789
+"""
+        self.assertTrue(repack_whl._has_all_quoted_filenames(record))
+
+    def test_none_quoted(self) -> None:
+        """Returns False when no lines are quoted (standard style)."""
+        record = """\
+torch/__init__.py,sha256=abc,123
+torch/utils.py,sha256=def,456
+torch-2.0.0.dist-info/WHEEL,sha256=ghi,789
+"""
+        self.assertFalse(repack_whl._has_all_quoted_filenames(record))
+
+    def test_mixed_quoting(self) -> None:
+        """Returns False when only some lines are quoted."""
+        record = """\
+"file,with,commas.py",sha256=abc,123
+normal_file.py,sha256=def,456
+"""
+        self.assertFalse(repack_whl._has_all_quoted_filenames(record))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/tools/wheelmaker_test.py

@@ -1,8 +1,45 @@
+import io
 import unittest
 
 import tools.wheelmaker as wheelmaker
 
 
+class QuoteAllFilenamesTest(unittest.TestCase):
+    """Tests for quote_all_filenames behavior in _WhlFile.
+
+    Some wheels (like torch) have all filenames quoted in their RECORD file.
+    When repacking, we preserve this style to minimize diffs.
+    """
+
+    def _make_whl_file(self, quote_all: bool) -> wheelmaker._WhlFile:
+        """Create a _WhlFile instance for testing."""
+        buf = io.BytesIO()
+        return wheelmaker._WhlFile(
+            buf,
+            mode="w",
+            distribution_prefix="test-1.0.0",
+            quote_all_filenames=quote_all,
+        )
+
+    def test_quote_all_quotes_simple_filenames(self) -> None:
+        """When quote_all_filenames=True, all filenames are quoted."""
+        whl = self._make_whl_file(quote_all=True)
+        self.assertEqual(whl._quote_filename("foo/bar.py"), '"foo/bar.py"')
+
+    def test_quote_all_false_leaves_simple_filenames_unquoted(self) -> None:
+        """When quote_all_filenames=False, simple filenames stay unquoted."""
+        whl = self._make_whl_file(quote_all=False)
+        self.assertEqual(whl._quote_filename("foo/bar.py"), "foo/bar.py")
+
+    def test_quote_all_quotes_filenames_with_commas(self) -> None:
+        """Filenames with commas are always quoted, regardless of quote_all_filenames."""
+        whl = self._make_whl_file(quote_all=True)
+        self.assertEqual(whl._quote_filename("foo,bar/baz.py"), '"foo,bar/baz.py"')
+
+        whl = self._make_whl_file(quote_all=False)
+        self.assertEqual(whl._quote_filename("foo,bar/baz.py"), '"foo,bar/baz.py"')
+
+
 class ArcNameFromTest(unittest.TestCase):
     def test_arcname_from(self) -> None:
         # (name, distribution_prefix, strip_path_prefixes, want) tuples