Fix importlib.machinery.FrozenImporter.load_module() to set __package__

properly. Discovered by also moving the loader tests over to
importlib.test.abc.LoaderTests.

3 files changed, 57 insertions, 20 deletions

diff --git a/Lib/importlib/NOTES b/Lib/importlib/NOTES
index 9d5fd69..3827713 100644
--- a/Lib/importlib/NOTES
+++ b/Lib/importlib/NOTES
@@ -3,7 +3,6 @@ to do
 
 * Use test.abc.LoaderTests
 
-    + frozen
     + source
 
 * Reorganize support code.
diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
index c1d09dd..d2d5d34 100644
--- a/Lib/importlib/_bootstrap.py
+++ b/Lib/importlib/_bootstrap.py
@@ -137,7 +137,12 @@ class FrozenImporter:
         """Load a frozen module."""
         if cls.find_module(fullname) is None:
             raise ImportError("{0} is not a frozen module".format(fullname))
-        return imp.init_frozen(fullname)
+        module = imp.init_frozen(fullname)
+        if hasattr(module, '__path__'):
+            module.__package__ = module.__name__
+        elif '.' in module.__name__:
+            module.__package__ = module.__name__.rsplit('.', 1)[0]
+        return module
 
 
 class ChainedImporter(object):
diff --git a/Lib/importlib/test/frozen/test_loader.py b/Lib/importlib/test/frozen/test_loader.py
index 95854f2..63a9742 100644
--- a/Lib/importlib/test/frozen/test_loader.py
+++ b/Lib/importlib/test/frozen/test_loader.py
@@ -1,26 +1,59 @@
 from importlib import machinery
-from ..builtin import test_loader
-
-
-class LoaderTests(test_loader.LoaderTests):
-
-    name = '__phello__'
-    load_module = staticmethod(lambda name:
-                                machinery.FrozenImporter.load_module(name))
-    verification = {'__name__': '__phello__', '__file__': '<frozen>',
-                    '__package__': None, '__path__': ['__phello__']}
-
-
-class SubmoduleLoaderTests(LoaderTests):
-
-    name = '__phello__.spam'
-    verification = {'__name__': '__phello__.spam', '__file__': '<frozen>',
-                    '__package__': None}
+from .. import abc
+from .. import support
+
+
+class LoaderTests(abc.LoaderTests):
+
+    def test_module(self):
+        with support.uncache('__hello__'):
+            module = machinery.FrozenImporter.load_module('__hello__')
+            check = {'__name__': '__hello__', '__file__': '<frozen>',
+                        '__package__': None}
+            for attr, value in check.items():
+                self.assertEqual(getattr(module, attr), value)
+
+    def test_package(self):
+        with support.uncache('__phello__'):
+            module = machinery.FrozenImporter.load_module('__phello__')
+            check = {'__name__': '__phello__', '__file__': '<frozen>',
+                     '__package__': '__phello__', '__path__': ['__phello__']}
+            for attr, value in check.items():
+                attr_value = getattr(module, attr)
+                self.assertEqual(attr_value, value,
+                                 "for __phello__.%s, %r != %r" %
+                                 (attr, attr_value, value))
+
+    def test_lacking_parent(self):
+        with support.uncache('__phello__', '__phello__.spam'):
+            module = machinery.FrozenImporter.load_module('__phello__.spam')
+            check = {'__name__': '__phello__.spam', '__file__': '<frozen>',
+                     '__package__': '__phello__'}
+            for attr, value in check.items():
+                attr_value = getattr(module, attr)
+                self.assertEqual(attr_value, value,
+                                 "for __phello__.spam.%s, %r != %r" %
+                                 (attr, attr_value, value))
+
+    def test_module_reuse(self):
+        with support.uncache('__hello__'):
+            module1 = machinery.FrozenImporter.load_module('__hello__')
+            module2 = machinery.FrozenImporter.load_module('__hello__')
+            self.assert_(module1 is module2)
+
+    def test_state_after_failure(self):
+        # No way to trigger an error in a frozen module.
+        pass
+
+    def test_unloadable(self):
+        assert machinery.FrozenImporter.find_module('_not_real') is None
+        self.assertRaises(ImportError, machinery.FrozenImporter.load_module,
+                            '_not_real')
 
 
 def test_main():
     from test.support import run_unittest
-    run_unittest(LoaderTests, SubmoduleLoaderTests)
+    run_unittest(LoaderTests)
 
 
 if __name__ == '__main__':