gh-76351: Support custom extractor for ZipFile.extract and ZipFile.extractall

Doc/library/zipfile.rst

@@ -373,13 +373,14 @@ ZipFile objects
       object was changed from ``'r'`` to ``'rb'``.
 
 
-.. method:: ZipFile.extract(member, path=None, pwd=None)
+.. method:: ZipFile.extract(member, path=None, pwd=None, extractor=None)
 
    Extract a member from the archive to the current working directory; *member*
-   must be its full name or a :class:`ZipInfo` object.  Its file information is
-   extracted as accurately as possible.  *path* specifies a different directory
-   to extract to.  *member* can be a filename or a :class:`ZipInfo` object.
-   *pwd* is the password used for encrypted files as a :class:`bytes` object.
+   must be its full name or a :class:`ZipInfo` object.  *path* specifies a
+   different directory to extract to.  *pwd* is the password used for encrypted
+   files as a :class:`bytes` object.  *extractor* is a custom
+   :ref:`extractor class <extractors>` that handles file attribute restoration
+   (defaults to :class:`ZipExtractorBase`, which restores no attributes).
 
    Returns the normalized path created (a directory or new file).
 
@@ -393,20 +394,31 @@ ZipFile objects
       characters (``:``, ``<``, ``>``, ``|``, ``"``, ``?``, and ``*``)
       replaced by underscore (``_``).
 
+   .. note::
+
+      Use :meth:`extractall` when extracting multiple members.  The
+      :meth:`extract` method may not restore all file attributes for the whole
+      directory hierarchy reliably.
+
    .. versionchanged:: 3.6
       Calling :meth:`extract` on a closed ZipFile will raise a
       :exc:`ValueError`.  Previously, a :exc:`RuntimeError` was raised.
 
    .. versionchanged:: 3.6.2
       The *path* parameter accepts a :term:`path-like object`.
 
+   .. versionchanged:: next
+      Added the *extractor* parameter.
+
 
-.. method:: ZipFile.extractall(path=None, members=None, pwd=None)
+.. method:: ZipFile.extractall(path=None, members=None, pwd=None, extractor=None)
 
    Extract all members from the archive to the current working directory.  *path*
    specifies a different directory to extract to.  *members* is optional and must
    be a subset of the list returned by :meth:`namelist`.  *pwd* is the password
-   used for encrypted files as a :class:`bytes` object.
+   used for encrypted files as a :class:`bytes` object.  *extractor* is a custom
+   :ref:`extractor class <extractors>` that handles file attribute restoration
+   (defaults to :class:`ZipExtractorBase`, which restores no attributes).
 
    .. warning::
 
@@ -423,6 +435,9 @@ ZipFile objects
    .. versionchanged:: 3.6.2
       The *path* parameter accepts a :term:`path-like object`.
 
+   .. versionchanged:: next
+      Added the *extractor* parameter.
+
 
 .. method:: ZipFile.printdir()
 
@@ -1041,6 +1056,74 @@ Instances have the following methods and attributes:
    Size of the uncompressed file.
 
 
+.. _extractors:
+
+Extractors
+----------
+
+An extractor class can be passed to :meth:`ZipFile.extract` or
+:meth:`ZipFile.extractall` to handle file attribute restoration.
+
+To create a custom extractor, subclass :class:`ZipExtractorBase` and implement
+the :meth:`~ZipExtractorBase.restore_attributes` method.  Here is an example
+that ignores errors when restoring file attributes:
+
+.. code-block:: python
+
+   class SafeZipExtractorTimeMode(zipfile.ZipExtractorTimeMode):
+       def restore_attributes(self, targetpath, zinfo):
+           try:
+               super().restore_attributes(targetpath, zinfo)
+           except (OSError, OverflowError, ValueError):
+               pass
+
+   with zipfile.ZipFile('spam.zip') as myzip:
+       myzip.extractall('eggs', extractor=SafeZipExtractorTimeMode)
+
+
+.. class:: ZipExtractorBase(archive, path=None, pwd=None)
+
+   The base extractor class that restores no file attributes.
+
+   .. method:: restore_attributes(targetpath, zinfo)
+
+      Called for every extracted member *zinfo* at *targetpath* to restore
+      its attributes.  This method does nothing by default.  Subclasses
+      should override this method to customize attribute restoration.
+
+   .. method:: finalize(exc_type=None, exc_value=None, traceback=None)
+
+      Called after all members have been extracted (or when an exception
+      occurs) to restore attributes for extracted directories in reverse
+      (bottom-up) order.  This prevents issues where extracting contents into
+      a directory resets its modification time or fails if it's not writable.
+      Subclasses can override this method to customize cleanup or error
+      recovery logic.
+
+
+.. class:: ZipExtractorTime(archive, path=None, pwd=None)
+
+   A subclass of :class:`ZipExtractorBase` that restores file mtime and atime.
+
+
+.. class:: ZipExtractorTimeMode(archive, path=None, pwd=None)
+
+   A subclass of :class:`ZipExtractorTime` that restores file mtime, atime, and
+   all mode bits (``0o7777``).
+
+
+.. class:: ZipExtractorTimeModeSafe(archive, path=None, pwd=None)
+
+   A subclass of :class:`ZipExtractorTimeMode` that restores file mtime, atime,
+   and safe file mode bits (``0o777``).
+
+
+.. class:: ZipExtractorTimeModeX(archive, path=None, pwd=None)
+
+   A subclass of :class:`ZipExtractorTimeMode` that restores file mtime, atime,
+   and executable bits (``0o111``).
+
+
 .. _zipfile-commandline:
 .. program:: zipfile
 

Lib/test/test_zipfile/test_core.py

@@ -3948,6 +3948,144 @@ def _test_extract_hackers_arcnames(self, hacknames):
             unlink(TESTFN2)
 
 
+class ExtractorTests(unittest.TestCase):
+    test_data = [
+        ('folder1/', b'', (1990, 1, 2, 0, 0, 0), (0o47777 << 16) | 0x10),
+        ('folder1/file1.txt', b'qawsedrftg', (1990, 2, 2, 0, 0, 0), 0o106642 << 16),
+        ('implicit_folder/file2.txt', b'azsxdcfvgb', (1990, 3, 2, 0, 0, 0), 0o100246 << 16),
+    ]
+
+    @classmethod
+    def setUpClass(cls):
+        os.mkdir(TESTFNDIR)
+        cls.dmode = os.stat(TESTFNDIR).st_mode
+        rmtree(TESTFNDIR)
+
+        with open(TESTFN, 'wb'):
+            pass
+        cls.fmode = os.stat(TESTFN).st_mode
+        os.unlink(TESTFN)
+
+    def setUp(self):
+        os.mkdir(TESTFNDIR)
+
+    def tearDown(self):
+        os.unlink(TESTFN)
+        rmtree(TESTFNDIR)
+
+    def make_test_archive(self, data):
+        with zipfile.ZipFile(TESTFN, 'w') as zipfp:
+            for filename, content, dt, ext_attr in data:
+                zinfo = zipfile.ZipInfo(filename, dt)
+                zinfo.external_attr = ext_attr
+                zipfp.writestr(zinfo, content)
+
+    def test_extract_default(self):
+        """Should not restore attributes by default."""
+        self.make_test_archive(self.test_data)
+        with zipfile.ZipFile(TESTFN) as zipfp:
+            for filename, content, dt, ext_attr in self.test_data:
+                with self.subTest(filename=filename):
+                    zipfp.extract(filename, TESTFNDIR)
+                    now = time.time()
+                    outfile = os.path.join(TESTFNDIR, filename)
+                    mtime = os.path.getmtime(outfile)
+                    self.assertAlmostEqual(mtime, now, delta=5)
+
+    def test_extract_time(self):
+        """Should restore mtime/atime."""
+        for extractor_cls, mask in (
+            (zipfile.ZipExtractorTime, 0),
+            (zipfile.ZipExtractorTimeModeX, 0o111),
+            (zipfile.ZipExtractorTimeModeSafe, 0o777),
+            (zipfile.ZipExtractorTimeMode, 0o7777),
+        ):
+            rmtree(TESTFNDIR)
+            self.make_test_archive(self.test_data)
+            with zipfile.ZipFile(TESTFN) as zipfp:
+                for filename, content, dt, ext_attr in self.test_data:
+                    with self.subTest(filename=filename, extractor=extractor_cls):
+                        zipfp.extract(filename, TESTFNDIR, extractor=extractor_cls)
+                        outfile = os.path.join(TESTFNDIR, filename)
+                        mtuple = time.localtime(os.path.getmtime(outfile))[:6]
+                        atuple = time.localtime(os.path.getatime(outfile))[:6]
+                        self.assertEqual(mtuple, dt)
+                        self.assertEqual(atuple, dt)
+
+    def test_extract_all_default(self):
+        """Should not restore attributes by default."""
+        self.make_test_archive(self.test_data)
+        now = time.time()
+        with zipfile.ZipFile(TESTFN) as zipfp:
+            zipfp.extractall(TESTFNDIR)
+        for filename, content, dt, ext_attr in self.test_data:
+            with self.subTest(filename=filename):
+                outfile = os.path.join(TESTFNDIR, filename)
+                mtime = os.path.getmtime(outfile)
+                self.assertAlmostEqual(mtime, now, delta=5)
+
+    def test_extract_all_time(self):
+        """Should restore mtime/atime."""
+        for extractor_cls, mask in (
+            (zipfile.ZipExtractorTime, 0),
+            (zipfile.ZipExtractorTimeModeX, 0o111),
+            (zipfile.ZipExtractorTimeModeSafe, 0o777),
+            (zipfile.ZipExtractorTimeMode, 0o7777),
+        ):
+            rmtree(TESTFNDIR)
+            self.make_test_archive(self.test_data)
+            with zipfile.ZipFile(TESTFN) as zipfp:
+                zipfp.extractall(TESTFNDIR, extractor=extractor_cls)
+            for filename, content, dt, ext_attr in self.test_data:
+                with self.subTest(filename=filename, extractor=extractor_cls):
+                    outfile = os.path.join(TESTFNDIR, filename)
+                    mtuple = time.localtime(os.path.getmtime(outfile))[:6]
+                    atuple = time.localtime(os.path.getatime(outfile))[:6]
+                    self.assertEqual(mtuple, dt)
+                    self.assertEqual(atuple, dt)
+
+    @unittest.skipIf(sys.platform == 'win32' or sys.platform == 'wasi', 'Requires file permissions')
+    def test_extract_all_mode(self):
+        """Should restore (masked) mode."""
+        for extractor_cls, mask in (
+            (zipfile.ZipExtractorTime, 0),
+            (zipfile.ZipExtractorTimeModeX, 0o111),
+            (zipfile.ZipExtractorTimeModeSafe, 0o777),
+            (zipfile.ZipExtractorTimeMode, 0o7777),
+        ):
+            rmtree(TESTFNDIR)
+            self.make_test_archive(self.test_data)
+            with zipfile.ZipFile(TESTFN) as zipfp:
+                zipfp.extractall(TESTFNDIR, extractor=extractor_cls)
+            for filename, content, dt, ext_attr in self.test_data:
+                with self.subTest(filename=filename, extractor=extractor_cls):
+                    outfile = os.path.join(TESTFNDIR, filename)
+                    default_mode = (self.dmode if os.path.isdir(outfile) else self.fmode)
+                    expected_mode = (default_mode & ~mask | ((ext_attr >> 16) & mask))
+                    mode = os.stat(outfile).st_mode
+                    self.assertEqual(oct(mode), oct(expected_mode))
+
+    @unittest.skipIf(sys.platform == 'win32' or sys.platform == 'wasi', 'Requires file permissions')
+    def test_extract_all_hierarchy(self):
+        """Should safely extract contents into a non-writable directory."""
+        test_data = [
+            ('folder1/', b'', (1990, 1, 2, 0, 0, 0), (0o40100 << 16) | 0x10),
+            ('folder1/file1.txt', b'qawsedrftg', (1990, 2, 2, 0, 0, 0), 0o100000 << 16),
+            ('folder1/file2.txt', b'azsxdcfvgb', (1990, 3, 2, 0, 0, 0), 0o100777 << 16),
+        ]
+        mask = 0o7777
+        self.make_test_archive(self.test_data)
+        with zipfile.ZipFile(TESTFN) as zipfp:
+            zipfp.extractall(TESTFNDIR, extractor=zipfile.ZipExtractorTimeMode)
+        for filename, content, dt, ext_attr in self.test_data:
+            with self.subTest(filename=filename):
+                outfile = os.path.join(TESTFNDIR, filename)
+                default_mode = (self.dmode if os.path.isdir(outfile) else self.fmode)
+                expected_mode = (default_mode & ~mask | ((ext_attr >> 16) & mask))
+                mode = os.stat(outfile).st_mode
+                self.assertEqual(oct(mode), oct(expected_mode))
+
+
 class OverwriteTests(archiver_tests.OverwriteTests, unittest.TestCase):
     testdir = TESTFN