Issue #15294: Fix a regression in pkgutil.extend_path()'s handling of nested namespace packages.

pitrou · pitrou · commit b2dd880e0a9e · 2012-07-09T21:23:58.000+02:00
diff --git a/Lib/pkgutil.py b/Lib/pkgutil.py
@@ -513,20 +513,30 @@ def extend_path(path, name):
         # frozen package.  Return the path unchanged in that case.
         return path
 
-    pname = os.path.join(*name.split('.')) # Reconstitute as relative path
     sname_pkg = name + ".pkg"
 
     path = path[:] # Start with a copy of the existing path
 
-    for dir in sys.path:
+    parent_package, _, final_name = name.rpartition('.')
+    if parent_package:
+        try:
+            search_path = sys.modules[parent_package].__path__
+        except (KeyError, AttributeError):
+            # We can't do anything: find_loader() returns None when
+            # passed a dotted name.
+            return path
+    else:
+        search_path = sys.path
+
+    for dir in search_path:
         if not isinstance(dir, str):
             continue
 
         finder = get_importer(dir)
         if finder is not None:
             # Is this finder PEP 420 compliant?
             if hasattr(finder, 'find_loader'):
-                loader, portions = finder.find_loader(name)
+                loader, portions = finder.find_loader(final_name)
             else:
                 # No, no need to call it
                 loader = None
diff --git a/Lib/test/test_pkgutil.py b/Lib/test/test_pkgutil.py
@@ -1,4 +1,4 @@
-from test.support import run_unittest
+from test.support import run_unittest, unload
 import unittest
 import sys
 import imp
@@ -214,8 +214,50 @@ def test_mixed_namespace(self):
     # XXX: test .pkg files
 
 
+class NestedNamespacePackageTest(unittest.TestCase):
+
+    def setUp(self):
+        self.basedir = tempfile.mkdtemp()
+        self.old_path = sys.path[:]
+
+    def tearDown(self):
+        sys.path[:] = self.old_path
+        shutil.rmtree(self.basedir)
+
+    def create_module(self, name, contents):
+        base, final = name.rsplit('.', 1)
+        base_path = os.path.join(self.basedir, base.replace('.', os.path.sep))
+        os.makedirs(base_path, exist_ok=True)
+        with open(os.path.join(base_path, final + ".py"), 'w') as f:
+            f.write(contents)
+
+    def test_nested(self):
+        pkgutil_boilerplate = (
+            'import pkgutil; '
+            '__path__ = pkgutil.extend_path(__path__, __name__)')
+        self.create_module('a.pkg.__init__', pkgutil_boilerplate)
+        self.create_module('b.pkg.__init__', pkgutil_boilerplate)
+        self.create_module('a.pkg.subpkg.__init__', pkgutil_boilerplate)
+        self.create_module('b.pkg.subpkg.__init__', pkgutil_boilerplate)
+        self.create_module('a.pkg.subpkg.c', 'c = 1')
+        self.create_module('b.pkg.subpkg.d', 'd = 2')
+        sys.path.insert(0, os.path.join(self.basedir, 'a'))
+        sys.path.insert(0, os.path.join(self.basedir, 'b'))
+        import pkg
+        self.addCleanup(unload, 'pkg')
+        self.assertEqual(len(pkg.__path__), 2)
+        import pkg.subpkg
+        self.addCleanup(unload, 'pkg.subpkg')
+        self.assertEqual(len(pkg.subpkg.__path__), 2)
+        from pkg.subpkg.c import c
+        from pkg.subpkg.d import d
+        self.assertEqual(c, 1)
+        self.assertEqual(d, 2)
+
+
 def test_main():
-    run_unittest(PkgutilTests, PkgutilPEP302Tests, ExtendPathTests)
+    run_unittest(PkgutilTests, PkgutilPEP302Tests, ExtendPathTests,
+                 NestedNamespacePackageTest)
     # this is necessary if test is run repeated (like when finding leaks)
     import zipimport
     zipimport._zip_directory_cache.clear()
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -31,6 +31,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #15294: Fix a regression in pkgutil.extend_path()'s handling of
+  nested namespace packages.
+
 - Issue #15056: imp.cache_from_source() and source_from_cache() raise
   NotImplementedError when sys.implementation.cache_tag is set to None.
 
