Implement the more specific PEP 302 semantics for loaders and what happens upon

load failure in relation to reloads. Also expose
importlib.util.module_for_loader to handle all of the details of this along
with making sure all current loaders behave nicely.

4 files changed, 126 insertions, 62 deletions

diff --git a/Lib/importlib/NOTES b/Lib/importlib/NOTES
index bdbedec..4a47385 100644
--- a/Lib/importlib/NOTES
+++ b/Lib/importlib/NOTES
@@ -1,68 +1,66 @@
 to do
 /////
 
-* Change failed loading based on PEP 302 changes.
-
 * Refactor source/bytecode finder/loader code such that bytecode support is a
   subclass of source support (makes it nicer for VMs that don't use CPython
   bytecode).
 
 * Implement PEP 302 protocol for loaders (should just be a matter of testing).
 
-    + Built-in.
-    + Frozen.
-    + Extension.
     + Source/bytecode.
 
 * Public API left to expose (w/ docs!)
 
-  + abc
+    + abc
+
+        - Finder
 
-      - Finder
+            * find_module
 
-        * find_module
+        - Loader
 
-      - Loader
+            * load_module
 
-        * load_module
+        - (?) Importer(Finder, Loader)
+        - ResourceLoader(Loader)
 
-      - (?) Importer(Finder, Loader)
+            * get_data
 
-      - ResourceLoader(Loader)
+        - InspectLoader(Loader)
 
-        * get_data
+            * is_package
+            * get_code
+            * get_source
 
-      - InspectLoader(Loader)
+        - (?) SourceLoader(ResourceLoader)
 
-        * is_package
-        * get_code
-        * get_source
+            * source_path
+            * bytecode_path
+            * write_bytecode (not abstract)
 
-      - (?) SourceLoader(ResourceLoader)
+    + util
 
-        * source_path
-        * bytecode_path
-        * write_bytecode (not abstract)
+        - set___package__ decorator
 
-  + util
+    + machinery
 
-      - get_module decorator (rename: module_for_loader)
-      - set___package__ decorator
+        - (?) Chained path hook/finder
+        - Extensions importers
 
-  + machinery
+            * ExtensionFinder
+            * (?) Loader
 
-      - (?) Chained path hook/finder
-      - Extensions importers
+        - Source/bytecode importers
 
-          * ExtensionFinder
-          * (?) Loader
+            * SourceFinder
+            * (?) Loader
 
-      - Source/bytecode importers
+    + test
 
-          * SourceFinder
-          * (?) Loader
+        - abc
 
-      - PathFinder
+            * FinderTests [doc]
+            * LoaderTests [doc]
 
 * Make sure that there is documentation *somewhere* fully explaining the
 semantics of import that can be referenced from the package's documentation
diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
index 1b4053d..65e9f17 100644
--- a/Lib/importlib/_bootstrap.py
+++ b/Lib/importlib/_bootstrap.py
@@ -136,8 +136,13 @@ class BuiltinImporter:
         """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)
-        return module
+        is_reload = fullname in sys.modules
+        try:
+            return imp.init_builtin(fullname)
+        except:
+            if not is_reload and fullname in sys.modules:
+                del sys.modules[fullname]
+            raise
 
 
 class FrozenImporter:
@@ -160,8 +165,13 @@ class FrozenImporter:
         """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)
-        return module
+        is_reload = fullname in sys.modules
+        try:
+            return imp.init_frozen(fullname)
+        except:
+            if not is_reload and fullname in sys.modules:
+                del sys.modules[fullname]
+            raise
 
 
 class ChainedImporter(object):
@@ -249,14 +259,13 @@ class _ExtensionFileLoader(object):
     @set___package__
     def load_module(self, fullname):
         """Load an extension module."""
-        assert self._name == fullname
+        is_reload = fullname in sys.modules
         try:
             module = imp.load_dynamic(fullname, self._path)
             module.__loader__ = self
             return module
         except:
-            # If an error occurred, don't leave a partially initialized module.
-            if fullname in sys.modules:
+            if not is_reload and fullname in sys.modules:
                 del sys.modules[fullname]
             raise
 
@@ -282,16 +291,17 @@ def suffix_list(suffix_type):
             if suffix[2] == suffix_type]
 
 
-# XXX Need a better name.
-def get_module(fxn):
-    """Decorator to handle selecting the proper module for load_module
-    implementations.
+def module_for_loader(fxn):
+    """Decorator to handle selecting the proper module for loaders.
 
     Decorated modules are passed the module to use instead of the module name.
     The module is either from sys.modules if it already exists (for reloading)
     or is a new module which has __name__ set. If any exception is raised by
-    the decorated method then __loader__, __name__, __file__, and __path__ are
-    all restored on the module to their original values.
+    the decorated method and the decorator added a module to sys.modules, then
+    the module is deleted from sys.modules.
+
+    The decorator assumes that the decorated method takes self/cls as a first
+    argument and the module as the second argument.
 
     """
     def decorated(self, fullname):
@@ -302,27 +312,12 @@ def get_module(fxn):
             # implicitly imports 'locale' and would otherwise trigger an
             # infinite loop.
             module = imp.new_module(fullname)
-            module.__name__ = fullname
             sys.modules[fullname] = module
-        else:
-            original_values = {}
-            modified_attrs = ['__loader__', '__name__', '__file__', '__path__']
-            for attr in modified_attrs:
-                try:
-                    original_values[attr] = getattr(module, attr)
-                except AttributeError:
-                    pass
         try:
             return fxn(self, module)
         except:
             if not is_reload:
                 del sys.modules[fullname]
-            else:
-                for attr in modified_attrs:
-                    if attr in original_values:
-                        setattr(module, attr, original_values[attr])
-                    elif hasattr(module, attr):
-                        delattr(module, attr)
             raise
     wrap(decorated, fxn)
     return decorated
@@ -375,7 +370,7 @@ class _PyFileLoader(object):
         return self._find_path(imp.PY_COMPILED)
 
     @check_name
-    @get_module
+    @module_for_loader
     def load_module(self, module):
         """Load a Python source or bytecode module."""
         name = module.__name__
diff --git a/Lib/importlib/test/test_util.py b/Lib/importlib/test/test_util.py
new file mode 100644
index 0000000..476b43f
--- /dev/null
+++ b/Lib/importlib/test/test_util.py
@@ -0,0 +1,69 @@
+from importlib import util
+from . import util as test_util
+import imp
+import sys
+import types
+import unittest
+
+
+class ModuleForLoaderTests(unittest.TestCase):
+
+    """Tests for importlib.util.module_for_loader."""
+
+    def return_module(self, name):
+        fxn = util.module_for_loader(lambda self, module: module)
+        return fxn(self, name)
+
+    def raise_exception(self, name):
+        def to_wrap(self, module):
+            raise ImportError
+        fxn = util.module_for_loader(to_wrap)
+        try:
+            fxn(self, name)
+        except ImportError:
+            pass
+
+    def test_new_module(self):
+        # Test that when no module exists in sys.modules a new module is
+        # created.
+        module_name = 'a.b.c'
+        with test_util.uncache(module_name):
+            module = self.return_module(module_name)
+            self.assert_(module_name in sys.modules)
+        self.assert_(isinstance(module, types.ModuleType))
+        self.assertEqual(module.__name__, module_name)
+
+    def test_reload(self):
+        # Test that a module is reused if already in sys.modules.
+        name = 'a.b.c'
+        module = imp.new_module('a.b.c')
+        with test_util.uncache(name):
+            sys.modules[name] = module
+            returned_module = self.return_module(name)
+            self.assert_(sys.modules[name] is returned_module)
+
+    def test_new_module_failure(self):
+        # Test that a module is removed from sys.modules if added but an
+        # exception is raised.
+        name = 'a.b.c'
+        with test_util.uncache(name):
+            self.raise_exception(name)
+            self.assert_(name not in sys.modules)
+
+    def test_reload_failure(self):
+        # Test that a failure on reload leaves the module in-place.
+        name = 'a.b.c'
+        module = imp.new_module(name)
+        with test_util.uncache(name):
+            sys.modules[name] = module
+            self.raise_exception(name)
+            self.assert_(sys.modules[name] is module)
+
+
+def test_main():
+    from test import support
+    support.run_unittest(ModuleForLoaderTests)
+
+
+if __name__ == '__main__':
+    test_main()
diff --git a/Lib/importlib/util.py b/Lib/importlib/util.py
new file mode 100644
index 0000000..6250d5b
--- /dev/null
+++ b/Lib/importlib/util.py
@@ -0,0 +1,2 @@
+"""Utility code for constructing importers, etc."""
+from ._bootstrap import module_for_loader