bpo-31638: Add compression support to zipapp (GH-3819)

Add optional argument `compressed` to `zipapp.create_archive`, and add
option `--compress` to the command line interface of `zipapp`.

5 files changed, 44 insertions, 7 deletions

diff --git a/Doc/library/zipapp.rst b/Doc/library/zipapp.rst
index 120bbbb..2bd65cb 100644
--- a/Doc/library/zipapp.rst
+++ b/Doc/library/zipapp.rst
@@ -79,6 +79,13 @@ The following options are understood:
 
    :option:`--main` cannot be specified when copying an archive.
 
+.. cmdoption:: -c, --compress
+
+   Compress files with the deflate method, reducing the size of the output
+   file. By default, files are stored uncompressed in the archive.
+
+   :option:`--compress` has no effect when copying an archive.
+
 .. cmdoption:: --info
 
    Display the interpreter embedded in the archive, for diagnostic purposes.  In
@@ -98,8 +105,7 @@ Python API
 The module defines two convenience functions:
 
 
-.. function:: create_archive(source, target=None, interpreter=None, main=None,
-                             filter=None)
+.. function:: create_archive(source, target=None, interpreter=None, main=None, filter=None, compressed=False)
 
    Create an application archive from *source*.  The source can be any
    of the following:
@@ -149,6 +155,11 @@ The module defines two convenience functions:
    (relative to the source directory).  It should return ``True`` if the
    file is to be added.
 
+   The optional *compressed* argument determines whether files are
+   compressed.  If set to ``True``, files in the archive are compressed
+   with the deflate method; otherwise, files are stored uncompressed.
+   This argument has no effect when copying an existing archive.
+
    If a file object is specified for *source* or *target*, it is the
    caller's responsibility to close it after calling create_archive.
 
@@ -159,7 +170,7 @@ The module defines two convenience functions:
    needed by that class.
 
    .. versionadded:: 3.7
-      Added the *filter* argument.
+      Added the *filter* and *compressed* arguments.
 
 .. function:: get_interpreter(archive)
 
diff --git a/Doc/whatsnew/3.7.rst b/Doc/whatsnew/3.7.rst
index a19a289..a474e76 100644
--- a/Doc/whatsnew/3.7.rst
+++ b/Doc/whatsnew/3.7.rst
@@ -285,9 +285,13 @@ zipapp
 ------
 
 Function :func:`zipapp.create_archive` now accepts an optional *filter*
-argument, to allow the user to select which files should be included in the
+argument to allow the user to select which files should be included in the
+archive, and an optional *compressed* argument to generate a compressed
 archive.
 
+A command line option ``--compress`` has also been added to support
+compression.
+
 
 Optimizations
 =============
diff --git a/Lib/test/test_zipapp.py b/Lib/test/test_zipapp.py
index 56cf37c..3602b99 100644
--- a/Lib/test/test_zipapp.py
+++ b/Lib/test/test_zipapp.py
@@ -100,6 +100,20 @@ class ZipAppTest(unittest.TestCase):
         expected_target = self.tmpdir / 'source.pyz'
         self.assertTrue(expected_target.is_file())
 
+    def test_create_archive_with_compression(self):
+        # Test packing a directory into a compressed archive.
+        source = self.tmpdir / 'source'
+        source.mkdir()
+        (source / '__main__.py').touch()
+        (source / 'test.py').touch()
+        target = self.tmpdir / 'source.pyz'
+
+        zipapp.create_archive(source, target, compressed=True)
+        with zipfile.ZipFile(target, 'r') as z:
+            for name in ('__main__.py', 'test.py'):
+                self.assertEqual(z.getinfo(name).compress_type,
+                                 zipfile.ZIP_DEFLATED)
+
     def test_no_main(self):
         # Test that packing a directory with no __main__.py fails.
         source = self.tmpdir / 'source'
diff --git a/Lib/zipapp.py b/Lib/zipapp.py
index 51d0290..ce77632 100644
--- a/Lib/zipapp.py
+++ b/Lib/zipapp.py
@@ -74,7 +74,7 @@ def _copy_archive(archive, new_archive, interpreter=None):
 
 
 def create_archive(source, target=None, interpreter=None, main=None,
-                   filter=None):
+                   filter=None, compressed=False):
     """Create an application archive from SOURCE.
 
     The SOURCE can be the name of a directory, or a filename or a file-like
@@ -133,7 +133,9 @@ def create_archive(source, target=None, interpreter=None, main=None,
 
     with _maybe_open(target, 'wb') as fd:
         _write_file_prefix(fd, interpreter)
-        with zipfile.ZipFile(fd, 'w') as z:
+        compression = (zipfile.ZIP_DEFLATED if compressed else
+                       zipfile.ZIP_STORED)
+        with zipfile.ZipFile(fd, 'w', compression=compression) as z:
             for child in source.rglob('*'):
                 arcname = child.relative_to(source)
                 if filter is None or filter(arcname):
@@ -170,6 +172,9 @@ def main(args=None):
     parser.add_argument('--main', '-m', default=None,
             help="The main function of the application "
                  "(default: use an existing __main__.py).")
+    parser.add_argument('--compress', '-c', action='store_true',
+            help="Compress files with the deflate method. "
+                 "Files are stored uncompressed by default.")
     parser.add_argument('--info', default=False, action='store_true',
             help="Display the interpreter from the archive.")
     parser.add_argument('source',
@@ -193,7 +198,8 @@ def main(args=None):
             raise SystemExit("Cannot change the main function when copying")
 
     create_archive(args.source, args.output,
-                   interpreter=args.python, main=args.main)
+                   interpreter=args.python, main=args.main,
+                   compressed=args.compress)
 
 
 if __name__ == '__main__':
diff --git a/Misc/NEWS.d/next/Library/2017-09-29-07-14-28.bpo-31638.jElfhl.rst b/Misc/NEWS.d/next/Library/2017-09-29-07-14-28.bpo-31638.jElfhl.rst
new file mode 100644
index 0000000..29e504e
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2017-09-29-07-14-28.bpo-31638.jElfhl.rst
@@ -0,0 +1,2 @@
+Add optional argument ``compressed`` to ``zipapp.create_archive``, and add
+option ``--compress`` to the command line interface of ``zipapp``.