Implement InspectLoader for FrozenImporter.

brettcannon · brettcannon · commit 8d11013169ef · 2009-03-15T02:20:16.000Z
diff --git a/Doc/library/importlib.rst b/Doc/library/importlib.rst
@@ -284,7 +284,8 @@ find and load modules.
 .. class:: FrozenImporter
 
     An :term:`importer` for frozen modules. This class implements the
-    :class:`importlib.abc.Finder` and :class:`importlib.abc.Loader` ABCs.
+    :class:`importlib.abc.Finder` and :class:`importlib.abc.InspectLoader`
+    ABCs.
 
     Only class methods are defined by this class to alleviate the need for
     instantiation.
diff --git 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
@@ -183,6 +183,16 @@ def wrapper(self, fullname):
     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 @@ def find_module(cls, fullname, path=None):
     @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 @@ def load_module(cls, fullname):
                 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
@@ -14,8 +14,6 @@ def load_module(self, fullname:str) -> types.ModuleType:
         """Abstract method which when implemented should load a module."""
         raise NotImplementedError
 
-Loader.register(machinery.FrozenImporter)
-
 
 class Finder(metaclass=abc.ABCMeta):
 
@@ -75,6 +73,7 @@ def get_source(self, fullname:str) -> str:
         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
@@ -1,4 +1,6 @@
 from importlib import machinery
+import imp
+import unittest
 from .. import abc
 from .. import util
 
@@ -53,9 +55,41 @@ def test_unloadable(self):
                             '_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__':
diff --git a/Python/import.c b/Python/import.c
@@ -1951,6 +1951,28 @@ get_frozen_object(char *name)
 	return PyMarshal_ReadObjectFromString((char *)p->code, size);
 }
 
+static PyObject *
+is_frozen_package(char *name)
+{
+	struct _frozen *p = find_frozen(name);
+	int size;
+
+	if (p == NULL) {
+		PyErr_Format(PyExc_ImportError,
+			     "No such frozen object named %.200s",
+			     name);
+		return NULL;
+	}
+
+	size = p->size;
+
+	if (size < 0)
+		Py_RETURN_TRUE;
+	else
+		Py_RETURN_FALSE;
+}
+
+
 /* Initialize a frozen module.
    Return 1 for success, 0 if the module is not found, and -1 with
    an exception set if the initialization failed.
@@ -2958,6 +2980,16 @@ imp_get_frozen_object(PyObject *self, PyObject *args)
 	return get_frozen_object(name);
 }
 
+static PyObject *
+imp_is_frozen_package(PyObject *self, PyObject *args)
+{
+	char *name;
+
+	if (!PyArg_ParseTuple(args, "s:is_frozen_package", &name))
+		return NULL;
+	return is_frozen_package(name);
+}
+
 static PyObject *
 imp_is_builtin(PyObject *self, PyObject *args)
 {
@@ -3191,6 +3223,7 @@ static PyMethodDef imp_methods[] = {
 	{"reload",       imp_reload,       METH_O,       doc_reload},
 	/* The rest are obsolete */
 	{"get_frozen_object",	imp_get_frozen_object,	METH_VARARGS},
+	{"is_frozen_package",   imp_is_frozen_package,  METH_VARARGS},
 	{"init_builtin",	imp_init_builtin,	METH_VARARGS},
 	{"init_frozen",		imp_init_frozen,	METH_VARARGS},
 	{"is_builtin",		imp_is_builtin,		METH_VARARGS},
