Support uncompressed PEXes.

pex/bin/pex.py

@@ -159,6 +159,21 @@ def configure_clp_pex_options(parser):
         ),
     )
 
+    group.add_argument(
+        "--compress",
+        "--compressed",
+        "--no-compress",
+        "--not-compressed",
+        "--no-compression",
+        dest="compress",
+        default=True,
+        action=HandleBoolAction,
+        help=(
+            "Whether to compress zip entries when creating either a zipapp PEX file or a packed "
+            "PEX's bootstrap and dependency zip files. Does nothing for loose layout PEXes."
+        ),
+    )
+
     runtime_mode = group.add_mutually_exclusive_group()
     runtime_mode.add_argument(
         "--unzip",
@@ -765,6 +780,7 @@ def do_main(
             bytecode_compile=options.compile,
             deterministic_timestamp=not options.use_system_time,
             layout=options.layout,
+            compress=options.compress,
         )
         if options.seed != Seed.NONE:
             seed_info = seed_cache(options, pex, verbose=options.seed == Seed.VERBOSE)

pex/common.py

@@ -745,6 +745,7 @@ def zip(
         exclude_file=lambda _: False,  # type: Callable[[str], bool]
         strip_prefix=None,  # type: Optional[str]
         labels=None,  # type: Optional[Iterable[str]]
+        compress=True,  # type: bool
     ):
         # type: (...) -> None
 
@@ -755,7 +756,8 @@ def zip(
         else:
             selected_files = self.files()
 
-        with open_zip(filename, mode) as zf:
+        compression = zipfile.ZIP_DEFLATED if compress else zipfile.ZIP_STORED
+        with open_zip(filename, mode, compression) as zf:
 
             def write_entry(
                 filename,  # type: str
@@ -769,7 +771,7 @@ def write_entry(
                     if deterministic_timestamp
                     else None,
                 )
-                zf.writestr(zip_entry.info, zip_entry.data)
+                zf.writestr(zip_entry.info, zip_entry.data, compression)
 
             def get_parent_dir(path):
                 # type: (str) -> Optional[str]

pex/pex_builder.py

@@ -600,6 +600,7 @@ def build(
         bytecode_compile=True,  # type: bool
         deterministic_timestamp=False,  # type: bool
         layout=Layout.ZIPAPP,  # type: Layout.Value
+        compress=True,  # type: bool
     ):
         # type: (...) -> None
         """Package the PEX application.
@@ -612,6 +613,8 @@ def build(
         :param bytecode_compile: If True, precompile .py files into .pyc files.
         :param deterministic_timestamp: If True, will use our hardcoded time for zipfile timestamps.
         :param layout: The layout to use for the PEX.
+        :param compress: Whether to compress zip entries when building to a layout that uses zip
+                         files.
 
         If the PEXBuilder is not yet frozen, it will be frozen by ``build``.  This renders the
         PEXBuilder immutable.
@@ -632,15 +635,20 @@ def build(
                     else:
                         os.mkdir(dirname)
                         self._build_packedapp(
-                            dirname=dirname, deterministic_timestamp=deterministic_timestamp
+                            dirname=dirname,
+                            deterministic_timestamp=deterministic_timestamp,
+                            compress=compress,
                         )
         else:
-            self._build_zipapp(filename=path, deterministic_timestamp=deterministic_timestamp)
+            self._build_zipapp(
+                filename=path, deterministic_timestamp=deterministic_timestamp, compress=compress
+            )
 
     def _build_packedapp(
         self,
         dirname,  # type: str
         deterministic_timestamp=False,  # type: bool
+        compress=True,  # type: bool
     ):
         # type: (...) -> None
 
@@ -654,13 +662,26 @@ def _build_packedapp(
                 safe_mkdir(os.path.dirname(dest))
                 safe_copy(os.path.realpath(os.path.join(self._chroot.chroot, f)), dest)
 
+        # Pex historically only supported compressed zips in packed layout, so we don't disturb the
+        # old cache structure for those zips and instead just use a subdir for un-compressed zips.
+        # This works for our two zip caches (we'll have no collisions with legacy compressed zips)
+        # since the bootstrap zip has a known name that is not "un-compressed" and "un-compressed"
+        # is not a valid wheel name either.
+        def zip_cache_dir(path):
+            # type: (str) -> str
+            if compress:
+                return path
+            return os.path.join(path, "un-compressed")
+
         # Zip up the bootstrap which is constant for a given version of Pex.
         bootstrap_hash = pex_info.bootstrap_hash
         if bootstrap_hash is None:
             raise AssertionError(
                 "Expected bootstrap_hash to be populated for {}.".format(self._pex_info)
             )
-        cached_bootstrap_zip_dir = os.path.join(pex_info.pex_root, "bootstrap_zips", bootstrap_hash)
+        cached_bootstrap_zip_dir = zip_cache_dir(
+            os.path.join(pex_info.pex_root, "bootstrap_zips", bootstrap_hash)
+        )
         with atomic_directory(
             cached_bootstrap_zip_dir, exclusive=False
         ) as atomic_bootstrap_zip_dir:
@@ -671,6 +692,7 @@ def _build_packedapp(
                     exclude_file=is_pyc_temporary_file,
                     strip_prefix=pex_info.bootstrap,
                     labels=("bootstrap",),
+                    compress=compress,
                 )
         safe_copy(
             os.path.join(cached_bootstrap_zip_dir, pex_info.bootstrap),
@@ -683,8 +705,8 @@ def _build_packedapp(
             internal_cache = os.path.join(dirname, pex_info.internal_cache)
             os.mkdir(internal_cache)
             for location, fingerprint in pex_info.distributions.items():
-                cached_installed_wheel_zip_dir = os.path.join(
-                    pex_info.pex_root, "installed_wheel_zips", fingerprint
+                cached_installed_wheel_zip_dir = zip_cache_dir(
+                    os.path.join(pex_info.pex_root, "installed_wheel_zips", fingerprint)
                 )
                 with atomic_directory(
                     cached_installed_wheel_zip_dir, exclusive=False
@@ -696,6 +718,7 @@ def _build_packedapp(
                             exclude_file=is_pyc_temporary_file,
                             strip_prefix=os.path.join(pex_info.internal_cache, location),
                             labels=(location,),
+                            compress=compress,
                         )
                 safe_copy(
                     os.path.join(cached_installed_wheel_zip_dir, location),
@@ -706,6 +729,7 @@ def _build_zipapp(
         self,
         filename,  # type: str
         deterministic_timestamp=False,  # type: bool
+        compress=True,  # type: bool
     ):
         # type: (...) -> None
         tmp_zip = filename + "~"
@@ -733,6 +757,7 @@ def _build_zipapp(
                 # pyc files that we should avoid copying since they are unuseful and inherently
                 # racy.
                 exclude_file=is_pyc_temporary_file,
+                compress=compress,
             )
         if os.path.exists(filename):
             os.unlink(filename)

tests/test_pex_builder.py

@@ -16,7 +16,7 @@
 from pex.layout import Layout
 from pex.pex import PEX
 from pex.pex_builder import CopyMode, PEXBuilder
-from pex.testing import PY_VER, built_wheel, make_bdist, make_env
+from pex.testing import PY_VER, WheelBuilder, built_wheel, make_bdist, make_env
 from pex.testing import write_simple_pex as write_pex
 from pex.typing import TYPE_CHECKING
 from pex.variables import ENV
@@ -495,3 +495,49 @@ def assert_pex_env_var_nested(**env):
     assert_pex_env_var_nested()
     assert_pex_env_var_nested(PEX_VENV=False)
     assert_pex_env_var_nested(PEX_VENV=True)
+
+
+@pytest.mark.parametrize(
+    "layout", [pytest.param(layout, id=layout.value) for layout in (Layout.PACKED, Layout.ZIPAPP)]
+)
+def test_build_compression(
+    tmpdir,  # type: Any
+    layout,  # type: Layout.Value
+    pex_project_dir,  # type: str
+):
+    # type: (...) -> None
+
+    pex_root = os.path.join(str(tmpdir), "pex_root")
+
+    pb = PEXBuilder()
+    pb.info.pex_root = pex_root
+    with ENV.patch(PEX_ROOT=pex_root):
+        pb.add_dist_location(WheelBuilder(pex_project_dir).bdist())
+    exe = os.path.join(str(tmpdir), "exe.py")
+    with open(exe, "w") as fp:
+        fp.write("import pex; print(pex.__file__)")
+    pb.set_executable(exe)
+
+    def assert_pex(pex):
+        # type: (str) -> None
+        assert (
+            subprocess.check_output(args=[sys.executable, pex]).decode("utf-8").startswith(pex_root)
+        )
+
+    compressed_pex = os.path.join(str(tmpdir), "compressed.pex")
+    pb.build(compressed_pex, layout=layout)
+    assert_pex(compressed_pex)
+
+    uncompressed_pex = os.path.join(str(tmpdir), "uncompressed.pex")
+    pb.build(uncompressed_pex, layout=layout, compress=False)
+    assert_pex(uncompressed_pex)
+
+    def size(pex):
+        # type: (str) -> int
+        if os.path.isfile(pex):
+            return os.path.getsize(pex)
+        return sum(
+            os.path.getsize(os.path.join(root, f)) for root, _, files in os.walk(pex) for f in files
+        )
+
+    assert size(compressed_pex) < size(uncompressed_pex)