Issue #14209: pkgutil.iter_zipimport_modules ignores the prefix for packages

ambv · ambv · commit 0d18c15fbf43 · 2016-06-11T18:02:46.000-07:00
Patch by James Pickering.
diff --git a/Lib/pkgutil.py b/Lib/pkgutil.py
@@ -375,7 +375,7 @@ def iter_zipimport_modules(importer, prefix=''):
             if len(fn)==2 and fn[1].startswith('__init__.py'):
                 if fn[0] not in yielded:
                     yielded[fn[0]] = 1
-                    yield fn[0], True
+                    yield prefix + fn[0], True
 
             if len(fn)!=1:
                 continue
diff --git a/Lib/test/test_pkgutil.py b/Lib/test/test_pkgutil.py
@@ -7,7 +7,6 @@
 import os
 import os.path
 import tempfile
-import types
 import shutil
 import zipfile
 
@@ -101,6 +100,83 @@ def test_unreadable_dir_on_syspath(self):
         for t in pkgutil.walk_packages(path=[self.dirname]):
             self.fail("unexpected package found")
 
+    def test_walkpackages_filesys(self):
+        pkg1 = 'test_walkpackages_filesys'
+        pkg1_dir = os.path.join(self.dirname, pkg1)
+        os.mkdir(pkg1_dir)
+        f = open(os.path.join(pkg1_dir, '__init__.py'), "wb")
+        f.close()
+        os.mkdir(os.path.join(pkg1_dir, 'sub'))
+        f = open(os.path.join(pkg1_dir, 'sub', '__init__.py'), "wb")
+        f.close()
+        f = open(os.path.join(pkg1_dir, 'sub', 'mod.py'), "wb")
+        f.close()
+
+        # Now, to juice it up, let's add the opposite packages, too.
+        pkg2 = 'sub'
+        pkg2_dir = os.path.join(self.dirname, pkg2)
+        os.mkdir(pkg2_dir)
+        f = open(os.path.join(pkg2_dir, '__init__.py'), "wb")
+        f.close()
+        os.mkdir(os.path.join(pkg2_dir, 'test_walkpackages_filesys'))
+        f = open(os.path.join(pkg2_dir, 'test_walkpackages_filesys', '__init__.py'), "wb")
+        f.close()
+        f = open(os.path.join(pkg2_dir, 'test_walkpackages_filesys', 'mod.py'), "wb")
+        f.close()
+
+        expected = [
+            'sub',
+            'sub.test_walkpackages_filesys',
+            'sub.test_walkpackages_filesys.mod',
+            'test_walkpackages_filesys',
+            'test_walkpackages_filesys.sub',
+            'test_walkpackages_filesys.sub.mod',
+        ]
+        actual= [e[1] for e in pkgutil.walk_packages([self.dirname])]
+        self.assertEqual(actual, expected)
+
+        for pkg in expected:
+            if pkg.endswith('mod'):
+                continue
+            del sys.modules[pkg]
+
+    def test_walkpackages_zipfile(self):
+        """Tests the same as test_walkpackages_filesys, only with a zip file."""
+
+        zip = 'test_walkpackages_zipfile.zip'
+        pkg1 = 'test_walkpackages_zipfile'
+        pkg2 = 'sub'
+
+        zip_file = os.path.join(self.dirname, zip)
+        z = zipfile.ZipFile(zip_file, 'w')
+        z.writestr(pkg2 + '/__init__.py', "")
+        z.writestr(pkg2 + '/' + pkg1 + '/__init__.py', "")
+        z.writestr(pkg2 + '/' + pkg1 + '/mod.py', "")
+        z.writestr(pkg1 + '/__init__.py', "")
+        z.writestr(pkg1 + '/' + pkg2 + '/__init__.py', "")
+        z.writestr(pkg1 + '/' + pkg2 + '/mod.py', "")
+        z.close()
+
+        sys.path.insert(0, zip_file)
+        expected = [
+            'sub',
+            'sub.test_walkpackages_zipfile',
+            'sub.test_walkpackages_zipfile.mod',
+            'test_walkpackages_zipfile',
+            'test_walkpackages_zipfile.sub',
+            'test_walkpackages_zipfile.sub.mod',
+        ]
+        actual= [e[1] for e in pkgutil.walk_packages([zip_file])]
+        self.assertEqual(actual, expected)
+        del sys.path[0]
+
+        for pkg in expected:
+            if pkg.endswith('mod'):
+                continue
+            del sys.modules[pkg]
+
+
+
 class PkgutilPEP302Tests(unittest.TestCase):
 
     class MyTestLoader(object):
@@ -324,11 +400,11 @@ def check_deprecated(self):
 
     def test_importer_deprecated(self):
         with self.check_deprecated():
-            x = pkgutil.ImpImporter("")
+            pkgutil.ImpImporter("")
 
     def test_loader_deprecated(self):
         with self.check_deprecated():
-            x = pkgutil.ImpLoader("", "", "", "")
+            pkgutil.ImpLoader("", "", "", "")
 
     def test_get_loader_avoids_emulation(self):
         with check_warnings() as w:
diff --git a/Misc/ACKS b/Misc/ACKS
@@ -1137,6 +1137,7 @@ Geoff Philbrick
 Gavrie Philipson
 Adrian Phillips
 Christopher J. Phoenix
+James Pickering
 Neale Pickett
 Jim St. Pierre
 Dan Pierson
