bpo-14678: Update zipimport to support importlib.invalidate_caches() (GH-24159)

desmondcheongzx · web-flow · commit 3abf6f010243 · 2021-03-08T12:06:02.000-08:00
Added an invalidate_caches() method to the zipimport.zipimporter class based on the implementation of importlib.FileFinder.invalidate_caches(). This was done by adding a get_files() method and an _archive_mtime attribute to zipimport.zipimporter to check for updates or cache invalidation whenever the cache of files and toc entry information in the zipimporter is accessed.
diff --git a/Doc/library/zipimport.rst b/Doc/library/zipimport.rst
@@ -166,6 +166,15 @@ zipimporter Objects
 
          Use :meth:`exec_module` instead.
 
+
+   .. method:: invalidate_caches()
+
+      Clear out the internal cache of information about files found within
+      the ZIP archive.
+
+      .. versionadded:: 3.10
+
+
    .. attribute:: archive
 
       The file name of the importer's associated ZIP file, without a possible
diff --git a/Lib/test/test_zipimport.py b/Lib/test/test_zipimport.py
@@ -506,6 +506,47 @@ def testZipImporterMethods(self):
         self.assertEqual(zi2.archive, TEMP_ZIP)
         self.assertEqual(zi2.prefix, TESTPACK + os.sep)
 
+    def testInvalidateCaches(self):
+        packdir = TESTPACK + os.sep
+        packdir2 = packdir + TESTPACK2 + os.sep
+        files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
+                 packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
+                 packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc),
+                 "spam" + pyc_ext: (NOW, test_pyc)}
+        self.addCleanup(os_helper.unlink, TEMP_ZIP)
+        with ZipFile(TEMP_ZIP, "w") as z:
+            for name, (mtime, data) in files.items():
+                zinfo = ZipInfo(name, time.localtime(mtime))
+                zinfo.compress_type = self.compression
+                zinfo.comment = b"spam"
+                z.writestr(zinfo, data)
+
+        zi = zipimport.zipimporter(TEMP_ZIP)
+        self.assertEqual(zi._files.keys(), files.keys())
+        # Check that the file information remains accurate after reloading
+        zi.invalidate_caches()
+        self.assertEqual(zi._files.keys(), files.keys())
+        # Add a new file to the ZIP archive
+        newfile = {"spam2" + pyc_ext: (NOW, test_pyc)}
+        files.update(newfile)
+        with ZipFile(TEMP_ZIP, "a") as z:
+            for name, (mtime, data) in newfile.items():
+                zinfo = ZipInfo(name, time.localtime(mtime))
+                zinfo.compress_type = self.compression
+                zinfo.comment = b"spam"
+                z.writestr(zinfo, data)
+        # Check that we can detect the new file after invalidating the cache
+        zi.invalidate_caches()
+        self.assertEqual(zi._files.keys(), files.keys())
+        spec = zi.find_spec('spam2')
+        self.assertIsNotNone(spec)
+        self.assertIsInstance(spec.loader, zipimport.zipimporter)
+        # Check that the cached data is removed if the file is deleted
+        os.remove(TEMP_ZIP)
+        zi.invalidate_caches()
+        self.assertIsNone(zi._files)
+        self.assertIsNone(zipimport._zip_directory_cache.get(zi.archive))
+
     def testZipImporterMethodsInSubDirectory(self):
         packdir = TESTPACK + os.sep
         packdir2 = packdir + TESTPACK2 + os.sep
diff --git a/Lib/zipimport.py b/Lib/zipimport.py
@@ -321,6 +321,16 @@ def get_resource_reader(self, fullname):
         return ZipReader(self, fullname)
 
 
+    def invalidate_caches(self):
+        """Reload the file data of the archive path."""
+        try:
+            self._files = _read_directory(self.archive)
+            _zip_directory_cache[self.archive] = self._files
+        except ZipImportError:
+            _zip_directory_cache.pop(self.archive, None)
+            self._files = None
+
+
     def __repr__(self):
         return f'<zipimporter object "{self.archive}{path_sep}{self.prefix}">'
 
diff --git a/Misc/NEWS.d/next/Library/2021-01-07-21-25-49.bpo-14678.1zniCH.rst b/Misc/NEWS.d/next/Library/2021-01-07-21-25-49.bpo-14678.1zniCH.rst
@@ -0,0 +1,3 @@
+Add an invalidate_caches() method to the zipimport.zipimporter class to
+support importlib.invalidate_caches().
+Patch by Desmond Cheong.
diff --git a/Python/importlib_zipimport.h b/Python/importlib_zipimport.h

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+Add an invalidate_caches() method to the zipimport.zipimporter class to`
	`2`	`+support importlib.invalidate_caches().`
	`3`	`+Patch by Desmond Cheong.`