Implement InspectLoader for FrozenImporter.

4 files changed, 65 insertions, 9 deletions

diff --git a/Lib/importlib/NOTES b/Lib/importlib/NOTES
index 44cfc09..b472fad 100644
--- a/Lib/importlib/NOTES
+++ b/Lib/importlib/NOTES
@@ -1,10 +1,6 @@
 to do
 /////
 
-* Implement InspectLoader for FrozenImporter.
-
-    + Expose function to see if a frozen module is a package.
-
 * Make sure that there is documentation *somewhere* fully explaining the
 semantics of import that can be referenced from the package's documentation
 (even if it is in the package documentation itself, although it might be best
diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
index 736e1b6..d5b909e 100644
--- a/Lib/importlib/_bootstrap.py
+++ b/Lib/importlib/_bootstrap.py
@@ -183,6 +183,16 @@ def _requires_builtin(fxn):
     return wrapper
 
 
+def _requires_frozen(fxn):
+    """Decorator to verify the named module is frozen."""
+    def wrapper(self, fullname):
+        if not imp.is_frozen(fullname):
+            raise ImportError("{0} is not a frozen module".format(fullname))
+        return fxn(self, fullname)
+    _wrap(wrapper, fxn)
+    return wrapper
+
+
 def _suffix_list(suffix_type):
     """Return a list of file suffixes based on the imp file type."""
     return [suffix[0] for suffix in imp.get_suffixes()
@@ -261,10 +271,9 @@ class FrozenImporter:
     @classmethod
     @set_package
     @set_loader
+    @_requires_frozen
     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))
         is_reload = fullname in sys.modules
         try:
             return imp.init_frozen(fullname)
@@ -273,6 +282,24 @@ class FrozenImporter:
                 del sys.modules[fullname]
             raise
 
+    @classmethod
+    @_requires_frozen
+    def get_code(cls, fullname):
+        """Return the code object for the frozen module."""
+        return imp.get_frozen_object(fullname)
+
+    @classmethod
+    @_requires_frozen
+    def get_source(cls, fullname):
+        """Return None as frozen modules do not have source code."""
+        return None
+
+    @classmethod
+    @_requires_frozen
+    def is_package(cls, fullname):
+        """Return if the frozen module is a package."""
+        return imp.is_frozen_package(fullname)
+
 
 class PyLoader:
 
diff --git a/Lib/importlib/abc.py b/Lib/importlib/abc.py
index c9c542a..b2bdb02 100644
--- a/Lib/importlib/abc.py
+++ b/Lib/importlib/abc.py
@@ -14,8 +14,6 @@ class Loader(metaclass=abc.ABCMeta):
         """Abstract method which when implemented should load a module."""
         raise NotImplementedError
 
-Loader.register(machinery.FrozenImporter)
-
 
 class Finder(metaclass=abc.ABCMeta):
 
@@ -75,6 +73,7 @@ class InspectLoader(Loader):
         return NotImplementedError
 
 InspectLoader.register(machinery.BuiltinImporter)
+InspectLoader.register(machinery.FrozenImporter)
 
 
 class PyLoader(_bootstrap.PyLoader, InspectLoader):
diff --git a/Lib/importlib/test/frozen/test_loader.py b/Lib/importlib/test/frozen/test_loader.py
index 54d805d..472d6ec 100644
--- a/Lib/importlib/test/frozen/test_loader.py
+++ b/Lib/importlib/test/frozen/test_loader.py
@@ -1,4 +1,6 @@
 from importlib import machinery
+import imp
+import unittest
 from .. import abc
 from .. import util
 
@@ -53,9 +55,41 @@ class LoaderTests(abc.LoaderTests):
                             '_not_real')
 
 
+class InspectLoaderTests(unittest.TestCase):
+
+    """Tests for the InspectLoader methods for FrozenImporter."""
+
+    def test_get_code(self):
+        # Make sure that the code object is good.
+        name = '__hello__'
+        code = machinery.FrozenImporter.get_code(name)
+        mod = imp.new_module(name)
+        exec(code, mod.__dict__)
+        self.assert_(hasattr(mod, 'initialized'))
+
+    def test_get_source(self):
+        # Should always return None.
+        result = machinery.FrozenImporter.get_source('__hello__')
+        self.assert_(result is None)
+
+    def test_is_package(self):
+        # Should be able to tell what is a package.
+        test_for = (('__hello__', False), ('__phello__', True),
+                    ('__phello__.spam', False))
+        for name, is_package in test_for:
+            result = machinery.FrozenImporter.is_package(name)
+            self.assert_(bool(result) == is_package)
+
+    def test_failure(self):
+        # Raise ImportError for modules that are not frozen.
+        for meth_name in ('get_code', 'get_source', 'is_package'):
+            method = getattr(machinery.FrozenImporter, meth_name)
+            self.assertRaises(ImportError, method, 'importlib')
+
+
 def test_main():
     from test.support import run_unittest
-    run_unittest(LoaderTests)
+    run_unittest(LoaderTests, InspectLoaderTests)
 
 
 if __name__ == '__main__':