Introduce importlib.util.ModuleManager which is a context manager to

handle providing (and cleaning up if needed) the module to be loaded.

A future commit will use the context manager in
Lib/importlib/_bootstrap.py and thus why the code is placed there
instead of in Lib/importlib/util.py.

Doc/library/importlib.rst

@@ -789,6 +789,15 @@ an :term:`importer`.
 
    .. versionadded:: 3.3
 
+.. class:: ModuleManager(name)
+
+    A :term:`context manager` which provides the module to load. The module will
+    either come from :attr:`sys.modules` in the case of reloading or a fresh
+    module if loading a new module. Proper cleanup of :attr:`sys.modules` occurs
+    if the module was new and an exception was raised.
+
+    .. versionadded:: 3.4
+
 .. decorator:: module_for_loader
 
     A :term:`decorator` for a :term:`loader` method,
@@ -818,6 +827,10 @@ an :term:`importer`.
     Use of this decorator handles all the details of which module object a
     loader should initialize as specified by :pep:`302` as best as possible.
 
+    .. note::
+       :class:`ModuleManager` subsumes the module management aspect of this
+       decorator.
+
     .. versionchanged:: 3.3
        :attr:`__loader__` and :attr:`__package__` are automatically set
        (when possible).

Lib/importlib/_bootstrap.py

@@ -9,7 +9,7 @@ work. One should use importlib as the public-facing version of this module.
 #
 # IMPORTANT: Whenever making changes to this module, be sure to run
 # a top-level make in order to get the frozen version of the module
-# update. Not doing so, will result in the Makefile to fail for
+# update. Not doing so will result in the Makefile to fail for
 # all others who don't have a ./python around to freeze the module
 # in the early stages of compilation.
 #
@@ -20,10 +20,6 @@ work. One should use importlib as the public-facing version of this module.
 # reference any injected objects! This includes not only global code but also
 # anything specified at the class level.
 
-# XXX Make sure all public names have no single leading underscore and all
-#     others do.
-
-
 # Bootstrap-related code ######################################################
 
 _CASE_INSENSITIVE_PLATFORMS = 'win', 'cygwin', 'darwin'
@@ -498,6 +494,38 @@ def _verbose_message(message, *args, verbosity=1):
         print(message.format(*args), file=sys.stderr)
 
 
+class ModuleManager:
+
+    """Context manager which returns the module to be loaded.
+
+    Does the proper unloading from sys.modules upon failure.
+
+    """
+
+    def __init__(self, name):
+        self._name = name
+
+    def __enter__(self):
+        self._module = sys.modules.get(self._name)
+        self._is_reload = self._module is not None
+        if not self._is_reload:
+            # This must be done before open() is called as the 'io' module
+            # implicitly imports 'locale' and would otherwise trigger an
+            # infinite loop.
+            self._module = new_module(self._name)
+            # This must be done before putting the module in sys.modules
+            # (otherwise an optimization shortcut in import.c becomes wrong)
+            self._module.__initializing__ = True
+            sys.modules[self._name] = self._module
+        return self._module
+
+    def __exit__(self, *args):
+        self._module.__initializing__ = False
+        del self._module
+        if any(arg is not None for arg in args) and not self._is_reload:
+            del sys.modules[self._name]
+
+
 def set_package(fxn):
     """Set __package__ on the returned module."""
     def set_package_wrapper(*args, **kwargs):

Lib/importlib/util.py

@@ -1,5 +1,6 @@
 """Utility code for constructing importers, etc."""
 
+from ._bootstrap import ModuleManager
 from ._bootstrap import module_for_loader
 from ._bootstrap import set_loader
 from ._bootstrap import set_package

Lib/test/test_importlib/test_util.py

@@ -2,10 +2,60 @@ from importlib import util
 from . import util as test_util
 import imp
 import sys
+from test import support
 import types
 import unittest
 
 
+class ModuleManagerTests(unittest.TestCase):
+
+    module_name = 'ModuleManagerTest_module'
+
+    def setUp(self):
+        support.unload(self.module_name)
+        self.addCleanup(support.unload, self.module_name)
+
+    def test_new_module(self):
+        # Test a new module is created, inserted into sys.modules, has
+        # __initializing__ set to True after entering the context manager,
+        # and __initializing__ set to False after exiting.
+        with util.ModuleManager(self.module_name) as module:
+            self.assertIn(self.module_name, sys.modules)
+            self.assertIs(sys.modules[self.module_name], module)
+            self.assertTrue(module.__initializing__)
+        self.assertFalse(module.__initializing__)
+
+    def test_new_module_failed(self):
+        # Test the module is removed from sys.modules.
+        try:
+            with util.ModuleManager(self.module_name) as module:
+                self.assertIn(self.module_name, sys.modules)
+                raise exception
+        except Exception:
+            self.assertNotIn(self.module_name, sys.modules)
+        else:
+            self.fail('importlib.util.ModuleManager swallowed an exception')
+
+    def test_reload(self):
+        # Test that the same module is in sys.modules.
+        created_module = imp.new_module(self.module_name)
+        sys.modules[self.module_name] = created_module
+        with util.ModuleManager(self.module_name) as module:
+            self.assertIs(module, created_module)
+
+    def test_reload_failed(self):
+        # Test that the module was left in sys.modules.
+        created_module = imp.new_module(self.module_name)
+        sys.modules[self.module_name] = created_module
+        try:
+            with util.ModuleManager(self.module_name) as module:
+                raise Exception
+        except Exception:
+            self.assertIn(self.module_name, sys.modules)
+        else:
+            self.fail('importlib.util.ModuleManager swallowed an exception')
+
+
 class ModuleForLoaderTests(unittest.TestCase):
 
     """Tests for importlib.util.module_for_loader."""

Misc/NEWS

@@ -96,6 +96,9 @@ Core and Builtins
 Library
 -------
 
+- Add importlib.util.ModuleManager as a context manager to provide the proper
+  module object to load.
+
 - Issue #18025: Fixed a segfault in io.BufferedIOBase.readinto() when raw
   stream's read() returns more bytes than requested.