add filtering of individual files to PyZipFile

changed output of debug messages to say "path" or "file"
extended test for filtering  certain files in a package
added test for filtering files in a python dir (no package)

diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py
index f57da5f94d1ba0927c12b3456dd3d8e2445551ba..7249b13d4c2b51e2cd3e175826790b4309d246d1 100644 (file)
--- a/Lib/test/test_zipfile.py
+++ b/Lib/test/test_zipfile.py
@@ -604,12 +604,21 @@ class PyZipFileTests(unittest.TestCase):
             reportStr = reportSIO.getvalue()
             self.assertTrue('SyntaxError' in reportStr)
 
-            # then check that the filter works
+            # then check that the filter works on the whole package
             with captured_stdout() as reportSIO:
                 zipfp.writepy(packagedir, filterfunc=lambda whatever: False)
             reportStr = reportSIO.getvalue()
             self.assertTrue('SyntaxError' not in reportStr)
 
+            # then check that the filter works on individual files
+            with captured_stdout() as reportSIO:
+                zipfp.writepy(packagedir, filterfunc=lambda fn:
+                                                     'bad' not in fn)
+            reportStr = reportSIO.getvalue()
+            if reportStr:
+                print(reportStr)
+            self.assertTrue('SyntaxError' not in reportStr)
+
     def test_write_with_optimization(self):
         import email
         packagedir = os.path.dirname(email.__file__)
@@ -649,6 +658,26 @@ class PyZipFileTests(unittest.TestCase):
         finally:
             shutil.rmtree(TESTFN2)
 
+    def test_write_python_directory_filtered(self):
+        os.mkdir(TESTFN2)
+        try:
+            with open(os.path.join(TESTFN2, "mod1.py"), "w") as fp:
+                fp.write("print(42)\n")
+
+            with open(os.path.join(TESTFN2, "mod2.py"), "w") as fp:
+                fp.write("print(42 * 42)\n")
+
+            with TemporaryFile() as t, zipfile.PyZipFile(t, "w") as zipfp:
+                zipfp.writepy(TESTFN2, filterfunc=lambda fn:
+                                                  not fn.endswith('mod2.py'))
+
+                names = zipfp.namelist()
+                self.assertCompiledIn('mod1.py', names)
+                self.assertNotIn('mod2.py', names)
+
+        finally:
+            shutil.rmtree(TESTFN2)
+
     def test_write_non_pyfile(self):
         with TemporaryFile() as t, zipfile.PyZipFile(t, "w") as zipfp:
             with open(TESTFN, 'w') as f:

diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index b3554fb669c0c2b0bf552fb7feb5df41bbe94c6a..ca2611dc79fb2a27a7ec8f5b345c62644aa3a098 100644 (file)
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -1582,7 +1582,8 @@ class PyZipFile(ZipFile):
         """
         if filterfunc and not filterfunc(pathname):
             if self.debug:
-                print('pathname "%s" skipped by filterfunc' % pathname)
+                label = 'path' if os.path.isdir(pathname) else 'file'
+                print('%s "%s" skipped by filterfunc' % (label, pathname))
             return
         dir, name = os.path.split(pathname)
         if os.path.isdir(pathname):
@@ -1611,6 +1612,10 @@ class PyZipFile(ZipFile):
                             self.writepy(path, basename,
                                          filterfunc=filterfunc)  # Recursive call
                     elif ext == ".py":
+                        if filterfunc and not filterfunc(path):
+                            if self.debug:
+                                print('file "%s" skipped by filterfunc' % path)
+                            continue
                         fname, arcname = self._get_codename(path[0:-3],
                                                             basename)
                         if self.debug:
@@ -1624,6 +1629,10 @@ class PyZipFile(ZipFile):
                     path = os.path.join(pathname, filename)
                     root, ext = os.path.splitext(filename)
                     if ext == ".py":
+                        if filterfunc and not filterfunc(path):
+                            if self.debug:
+                                print('file "%s" skipped by filterfunc' % path)
+                            continue
                         fname, arcname = self._get_codename(path[0:-3],
                                                             basename)
                         if self.debug: