Move importlib completely over to using rpartition and accepting the empty

string for top-level modules.

4 files changed, 22 insertions, 15 deletions

diff --git a/Lib/importlib/NOTES b/Lib/importlib/NOTES
index 05920ad..f0d8e43 100644
--- a/Lib/importlib/NOTES
+++ b/Lib/importlib/NOTES
@@ -1,10 +1,7 @@
 to do
 /////
 
-* Use rpartition for getting the package of a module.
-
-    + Make sure there is a test for the empty string as acceptable for
-      __package__.
+* Extract test_path_hooks constants into a util module for extension testing.
 
 * Implement PEP 302 protocol for loaders (should just be a matter of testing).
 
diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
index bc50036..810f793 100644
--- a/Lib/importlib/_bootstrap.py
+++ b/Lib/importlib/_bootstrap.py
@@ -90,6 +90,18 @@ class closing:
         self.obj.close()
 
 
+def set___package__(fxn):
+    """Set __package__ on the returned module."""
+    def wrapper(*args, **kwargs):
+        module = fxn(*args, **kwargs)
+        if not hasattr(module, '__package__') or module.__package__ is None:
+            module.__package__ = module.__name__
+            if not hasattr(module, '__path__'):
+                module.__package__ = module.__package__.rpartition('.')[0]
+        return module
+    return wrapper
+
+
 class BuiltinImporter:
 
     """Meta path loader for built-in modules.
@@ -111,12 +123,12 @@ class BuiltinImporter:
         return cls if imp.is_builtin(fullname) else None
 
     @classmethod
+    @set___package__
     def load_module(cls, fullname):
         """Load a built-in module."""
         if fullname not in sys.builtin_module_names:
             raise ImportError("{0} is not a built-in module".format(fullname))
         module = imp.init_builtin(fullname)
-        module.__package__ = ''
         return module
 
 
@@ -135,14 +147,12 @@ class FrozenImporter:
         return cls if imp.is_frozen(fullname) else None
 
     @classmethod
+    @set___package__
     def load_module(cls, fullname):
         """Load a frozen module."""
         if cls.find_module(fullname) is None:
             raise ImportError("{0} is not a frozen module".format(fullname))
         module = imp.init_frozen(fullname)
-        module.__package__ = module.__name__
-        if not hasattr(module, '__path__'):
-            module.__package__ = module.__package__.rpartition('.')[0]
         return module
 
 
@@ -230,6 +240,7 @@ class _ExtensionFileLoader(object):
             raise ValueError("extension modules cannot be packages")
 
     @check_name
+    @set___package__
     def load_module(self, fullname):
         """Load an extension module."""
         assert self._name == fullname
@@ -368,11 +379,9 @@ class _PyFileLoader(object):
         module.__loader__ = self
         if self._is_pkg:
             module.__path__  = [module.__file__.rsplit(path_sep, 1)[0]]
-            module.__package__ = module.__name__
-        elif '.' in module.__name__:
-            module.__package__ = module.__name__.rsplit('.', 1)[0]
-        else:
-            module.__package__ = None
+        module.__package__ = module.__name__
+        if not hasattr(module, '__path__'):
+            module.__package__ = module.__package__.rpartition('.')[0]
         exec(code_object, module.__dict__)
         return module
 
diff --git a/Lib/importlib/test/extension/test_loader.py b/Lib/importlib/test/extension/test_loader.py
index 1e43fdd..01f3426 100644
--- a/Lib/importlib/test/extension/test_loader.py
+++ b/Lib/importlib/test/extension/test_loader.py
@@ -21,7 +21,8 @@ class LoaderTests(abc.LoaderTests):
         with util.uncache(test_path_hook.NAME):
             module = self.load_module(test_path_hook.NAME)
             for attr, value in [('__name__', test_path_hook.NAME),
-                                ('__file__', test_path_hook.FILEPATH)]:
+                                ('__file__', test_path_hook.FILEPATH),
+                                ('__package__', '')]:
                 self.assertEqual(getattr(module, attr), value)
             self.assert_(test_path_hook.NAME in sys.modules)
 
diff --git a/Lib/importlib/test/source/test_loader.py b/Lib/importlib/test/source/test_loader.py
index e333b85..4ca9af1 100644
--- a/Lib/importlib/test/source/test_loader.py
+++ b/Lib/importlib/test/source/test_loader.py
@@ -23,7 +23,7 @@ class SimpleTest(unittest.TestCase):
             module = loader.load_module('_temp')
             self.assert_('_temp' in sys.modules)
             check = {'__name__': '_temp', '__file__': mapping['_temp'],
-                     '__package__': None}
+                     '__package__': ''}
             for attr, value in check.items():
                 self.assertEqual(getattr(module, attr), value)