Issue #26039: zipfile.ZipFile.open() can now be used to write data into a ZIP

file, as well as for extracting data.  Patch by Thomas Kluyver.

Doc/library/zipfile.rst

@@ -207,15 +207,15 @@ ZipFile Objects
 .. index::
    single: universal newlines; zipfile.ZipFile.open method
 
-.. method:: ZipFile.open(name, mode='r', pwd=None)
+.. method:: ZipFile.open(name, mode='r', pwd=None, force_zip64=False)
 
-   Extract a member from the archive as a file-like object (ZipExtFile). *name*
-   is the name of the file in the archive, or a :class:`ZipInfo` object. The
+   Access a member of the archive as a file-like object.  *name*
+   is the name of the file in the archive, or a :class:`ZipInfo` object.  The
    *mode* parameter, if included, must be one of the following: ``'r'`` (the
-   default), ``'U'``, or ``'rU'``. Choosing ``'U'`` or  ``'rU'`` will enable
-   :term:`universal newlines` support in the read-only object.  *pwd* is the
-   password used for encrypted files.  Calling  :meth:`.open` on a closed
-   ZipFile will raise a  :exc:`RuntimeError`.
+   default), ``'U'``, ``'rU'`` or ``'w'``.  Choosing ``'U'`` or  ``'rU'`` will
+   enable :term:`universal newlines` support in the read-only object.  *pwd* is
+   the password used to decrypt encrypted ZIP files.  Calling  :meth:`.open` on
+   a closed ZipFile will raise a  :exc:`RuntimeError`.
 
    :meth:`~ZipFile.open` is also a context manager and therefore supports the
    :keyword:`with` statement::
@@ -224,17 +224,23 @@ ZipFile Objects
           with myzip.open('eggs.txt') as myfile:
               print(myfile.read())
 
-   .. note::
-
-      The file-like object is read-only and provides the following methods:
-      :meth:`~io.BufferedIOBase.read`, :meth:`~io.IOBase.readline`,
-      :meth:`~io.IOBase.readlines`, :meth:`__iter__`,
-      :meth:`~iterator.__next__`.
+   With *mode* ``'r'``, ``'U'`` or ``'rU'``, the file-like object
+   (``ZipExtFile``) is read-only and provides the following methods:
+   :meth:`~io.BufferedIOBase.read`, :meth:`~io.IOBase.readline`,
+   :meth:`~io.IOBase.readlines`, :meth:`__iter__`,
+   :meth:`~iterator.__next__`.  These objects can operate independently of
+   the ZipFile.
 
-   .. note::
+   With ``mode='w'``, a writable file handle is returned, which supports the
+   :meth:`~io.BufferedIOBase.write` method.  While a writable file handle is open,
+   attempting to read or write other files in the ZIP file will raise a
+   :exc:`RuntimeError`.
 
-      Objects returned by :meth:`.open` can operate independently of the
-      ZipFile.
+   When writing a file, if the file size is not known in advance but may exceed
+   2 GiB, pass ``force_zip64=True`` to ensure that the header format is
+   capable of supporting large files.  If the file size is known in advance,
+   construct a :class:`ZipInfo` object with :attr:`~ZipInfo.file_size` set, and
+   use that as the *name* parameter.
 
    .. note::
 
@@ -246,6 +252,10 @@ ZipFile Objects
       The ``'U'`` or  ``'rU'`` mode.  Use :class:`io.TextIOWrapper` for reading
       compressed text files in :term:`universal newlines` mode.
 
+   .. versionchanged:: 3.6
+      :meth:`open` can now be used to write files into the archive with the
+      ``mode='w'`` option.
+
 .. method:: ZipFile.extract(member, path=None, pwd=None)
 
    Extract a member from the archive to the current working directory; *member*

Doc/whatsnew/3.6.rst

@@ -350,6 +350,10 @@ A new :meth:`ZipInfo.is_dir() <zipfile.ZipInfo.is_dir>` method can be used
 to check if the :class:`~zipfile.ZipInfo` instance represents a directory.
 (Contributed by Thomas Kluyver in :issue:`26039`.)
 
+The :meth:`ZipFile.open() <zipfile.ZipFile.open>` method can now be used to
+write data into a ZIP file, as well as for extracting data.
+(Contributed by Thomas Kluyver in :issue:`26039`.)
+
 
 zlib
 ----

Lib/test/test_zipfile.py

@@ -61,6 +61,9 @@ class AbstractTestsWithSourceFile:
             zipfp.write(TESTFN, "another.name")
             zipfp.write(TESTFN, TESTFN)
             zipfp.writestr("strfile", self.data)
+            with zipfp.open('written-open-w', mode='w') as f:
+                for line in self.line_gen:
+                    f.write(line)
 
     def zip_test(self, f, compression):
         self.make_test_archive(f, compression)
@@ -76,7 +79,7 @@ class AbstractTestsWithSourceFile:
             zipfp.printdir(file=fp)
             directory = fp.getvalue()
             lines = directory.splitlines()
-            self.assertEqual(len(lines), 4) # Number of files + header
+            self.assertEqual(len(lines), 5) # Number of files + header
 
             self.assertIn('File Name', lines[0])
             self.assertIn('Modified', lines[0])
@@ -90,23 +93,25 @@ class AbstractTestsWithSourceFile:
 
             # Check the namelist
             names = zipfp.namelist()
-            self.assertEqual(len(names), 3)
+            self.assertEqual(len(names), 4)
             self.assertIn(TESTFN, names)
             self.assertIn("another.name", names)
             self.assertIn("strfile", names)
+            self.assertIn("written-open-w", names)
 
             # Check infolist
             infos = zipfp.infolist()
             names = [i.filename for i in infos]
-            self.assertEqual(len(names), 3)
+            self.assertEqual(len(names), 4)
             self.assertIn(TESTFN, names)
             self.assertIn("another.name", names)
             self.assertIn("strfile", names)
+            self.assertIn("written-open-w", names)
             for i in infos:
                 self.assertEqual(i.file_size, len(self.data))
 
             # check getinfo
-            for nm in (TESTFN, "another.name", "strfile"):
+            for nm in (TESTFN, "another.name", "strfile", "written-open-w"):
                 info = zipfp.getinfo(nm)
                 self.assertEqual(info.filename, nm)
                 self.assertEqual(info.file_size, len(self.data))
@@ -372,14 +377,18 @@ class StoredTestsWithSourceFile(AbstractTestsWithSourceFile,
     test_low_compression = None
 
     def zip_test_writestr_permissions(self, f, compression):
-        # Make sure that writestr creates files with mode 0600,
-        # when it is passed a name rather than a ZipInfo instance.
+        # Make sure that writestr and open(... mode='w') create files with
+        # mode 0600, when they are passed a name rather than a ZipInfo
+        # instance.
 
         self.make_test_archive(f, compression)
         with zipfile.ZipFile(f, "r") as zipfp:
             zinfo = zipfp.getinfo('strfile')
             self.assertEqual(zinfo.external_attr, 0o600 << 16)
 
+            zinfo2 = zipfp.getinfo('written-open-w')
+            self.assertEqual(zinfo2.external_attr, 0o600 << 16)
+
     def test_writestr_permissions(self):
         for f in get_files(self):
             self.zip_test_writestr_permissions(f, zipfile.ZIP_STORED)
@@ -451,6 +460,10 @@ class StoredTestsWithSourceFile(AbstractTestsWithSourceFile,
         with zipfile.ZipFile(TESTFN2, mode="r") as zipfp:
             self.assertRaises(RuntimeError, zipfp.write, TESTFN)
 
+        with zipfile.ZipFile(TESTFN2, mode="r") as zipfp:
+            with self.assertRaises(RuntimeError):
+                zipfp.open(TESTFN, mode='w')
+
     def test_add_file_before_1980(self):
         # Set atime and mtime to 1970-01-01
         os.utime(TESTFN, (0, 0))
@@ -1428,6 +1441,35 @@ class OtherTests(unittest.TestCase):
             # testzip returns the name of the first corrupt file, or None
             self.assertIsNone(zipf.testzip())
 
+    def test_open_conflicting_handles(self):
+        # It's only possible to open one writable file handle at a time
+        msg1 = b"It's fun to charter an accountant!"
+        msg2 = b"And sail the wide accountant sea"
+        msg3 = b"To find, explore the funds offshore"
+        with zipfile.ZipFile(TESTFN2, 'w', zipfile.ZIP_STORED) as zipf:
+            with zipf.open('foo', mode='w') as w2:
+                w2.write(msg1)
+            with zipf.open('bar', mode='w') as w1:
+                with self.assertRaises(RuntimeError):
+                    zipf.open('handle', mode='w')
+                with self.assertRaises(RuntimeError):
+                    zipf.open('foo', mode='r')
+                with self.assertRaises(RuntimeError):
+                    zipf.writestr('str', 'abcde')
+                with self.assertRaises(RuntimeError):
+                    zipf.write(__file__, 'file')
+                with self.assertRaises(RuntimeError):
+                    zipf.close()
+                w1.write(msg2)
+            with zipf.open('baz', mode='w') as w2:
+                w2.write(msg3)
+
+        with zipfile.ZipFile(TESTFN2, 'r') as zipf:
+            self.assertEqual(zipf.read('foo'), msg1)
+            self.assertEqual(zipf.read('bar'), msg2)
+            self.assertEqual(zipf.read('baz'), msg3)
+            self.assertEqual(zipf.namelist(), ['foo', 'bar', 'baz'])
+
     def tearDown(self):
         unlink(TESTFN)
         unlink(TESTFN2)
@@ -1761,6 +1803,22 @@ class UnseekableTests(unittest.TestCase):
                     with zipf.open('twos') as zopen:
                         self.assertEqual(zopen.read(), b'222')
 
+    def test_open_write(self):
+        for wrapper in (lambda f: f), Tellable, Unseekable:
+            with self.subTest(wrapper=wrapper):
+                f = io.BytesIO()
+                f.write(b'abc')
+                bf = io.BufferedWriter(f)
+                with zipfile.ZipFile(wrapper(bf), 'w', zipfile.ZIP_STORED) as zipf:
+                    with zipf.open('ones', 'w') as zopen:
+                        zopen.write(b'111')
+                    with zipf.open('twos', 'w') as zopen:
+                        zopen.write(b'222')
+                self.assertEqual(f.getvalue()[:5], b'abcPK')
+                with zipfile.ZipFile(f) as zipf:
+                    self.assertEqual(zipf.read('ones'), b'111')
+                    self.assertEqual(zipf.read('twos'), b'222')
+
 
 @requires_zlib
 class TestsWithMultipleOpens(unittest.TestCase):
@@ -1870,6 +1928,19 @@ class TestsWithMultipleOpens(unittest.TestCase):
         with open(os.devnull) as f:
             self.assertLess(f.fileno(), 100)
 
+    def test_write_while_reading(self):
+        with zipfile.ZipFile(TESTFN2, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            zipf.writestr('ones', self.data1)
+        with zipfile.ZipFile(TESTFN2, 'a', zipfile.ZIP_DEFLATED) as zipf:
+            with zipf.open('ones', 'r') as r1:
+                data1 = r1.read(500)
+                with zipf.open('twos', 'w') as w1:
+                    w1.write(self.data2)
+                data1 += r1.read()
+        self.assertEqual(data1, self.data1)
+        with zipfile.ZipFile(TESTFN2) as zipf:
+            self.assertEqual(zipf.read('twos'), self.data2)
+
     def tearDown(self):
         unlink(TESTFN2)
 

Misc/NEWS

@@ -277,6 +277,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #26039: zipfile.ZipFile.open() can now be used to write data into a ZIP
+  file, as well as for extracting data.  Patch by Thomas Kluyver.
+
 - Issue #26892: Honor debuglevel flag in urllib.request.HTTPHandler. Patch
   contributed by Chi Hsuan Yen.