From b52f3818d8ef896c7b634c34676c90670e04c242 Mon Sep 17 00:00:00 2001
From: Keith Smiley <keithbsmiley@gmail.com>
Date: Tue, 1 Oct 2024 10:15:08 -0700
Subject: [PATCH] feat: add py_wheel.compress to control using compression
 (#2260)

This is useful for development where, for large wheels, a significant
portion of the time can be spent compression native binaries

Fixes https://github.com/bazelbuild/rules_python/issues/2240
---
 CHANGELOG.md                | 3 ++-
 python/private/py_wheel.bzl | 7 +++++++
 tools/wheelmaker.py         | 9 +++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8982a5f716..a16a5e3c81 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -31,7 +31,8 @@ A brief description of the categories of changes:
 * (whl_filegroup): Provide per default also the `RECORD` file
 
 ### Added
-* Nothing yet
+* (py_wheel) Now supports `compress = (True|False)` to allow disabling
+  compression to speed up development.
 
 ### Removed
 * Nothing yet
diff --git a/python/private/py_wheel.bzl b/python/private/py_wheel.bzl
index ef9e6f24ae..26eb5652a6 100644
--- a/python/private/py_wheel.bzl
+++ b/python/private/py_wheel.bzl
@@ -34,6 +34,10 @@ _distribution_attrs = {
         default = "none",
         doc = "Python ABI tag. 'none' for pure-Python wheels.",
     ),
+    "compress": attr.bool(
+        default = True,
+        doc = "Enable compression of the final archive.",
+    ),
     "distribution": attr.string(
         mandatory = True,
         doc = """\
@@ -466,6 +470,9 @@ def _py_wheel_impl(ctx):
         args.add("--description_file", description_file)
         other_inputs.append(description_file)
 
+    if not ctx.attr.compress:
+        args.add("--no_compress")
+
     for target, filename in ctx.attr.extra_distinfo_files.items():
         target_files = target.files.to_list()
         if len(target_files) != 1:
diff --git a/tools/wheelmaker.py b/tools/wheelmaker.py
index 68578b8e58..db287ebaee 100644
--- a/tools/wheelmaker.py
+++ b/tools/wheelmaker.py
@@ -227,6 +227,7 @@ def __init__(
         python_tag,
         abi,
         platform,
+        compress,
         outfile=None,
         strip_path_prefixes=None,
     ):
@@ -238,6 +239,7 @@ def __init__(
         self._platform = platform
         self._outfile = outfile
         self._strip_path_prefixes = strip_path_prefixes
+        self._compress = compress
         self._wheelname_fragment_distribution_name = escape_filename_distribution_name(
             self._name
         )
@@ -254,6 +256,7 @@ def __enter__(self):
             mode="w",
             distribution_prefix=self._distribution_prefix,
             strip_path_prefixes=self._strip_path_prefixes,
+            compression=zipfile.ZIP_DEFLATED if self._compress else zipfile.ZIP_STORED,
         )
         return self
 
@@ -388,6 +391,11 @@ def parse_args() -> argparse.Namespace:
     output_group.add_argument(
         "--out", type=str, default=None, help="Override name of ouptut file"
     )
+    output_group.add_argument(
+        "--no_compress",
+        action="store_true",
+        help="Disable compression of the final archive",
+    )
     output_group.add_argument(
         "--name_file",
         type=Path,
@@ -516,6 +524,7 @@ def main() -> None:
         platform=arguments.platform,
         outfile=arguments.out,
         strip_path_prefixes=strip_prefixes,
+        compress=not arguments.no_compress,
     ) as maker:
         for package_filename, real_filename in all_files:
             maker.add_file(package_filename, real_filename)