From 5533ff6a2ea1560cf6700c07096f34c5b9bf5874 Mon Sep 17 00:00:00 2001
From: Nick Coghlan <ncoghlan@gmail.com>
Date: Wed, 22 Apr 2009 15:26:04 +0000
Subject: [PATCH] Issue 5354: Change API for import_fresh_module() to better
 support test_warnings use case (also fixes some bugs in the original
 implementation)

---
 Doc/library/test.rst      | 32 +++++++++++++++++---
 Lib/test/test_heapq.py    |  2 +-
 Lib/test/test_support.py  | 62 ++++++++++++++++++++++++++++++---------
 Lib/test/test_warnings.py |  4 +--
 Misc/NEWS                 |  6 ++++
 5 files changed, 85 insertions(+), 21 deletions(-)

diff --git a/Doc/library/test.rst b/Doc/library/test.rst
index ccca2991c51..ba2c3b873d9 100644
--- a/Doc/library/test.rst
+++ b/Doc/library/test.rst
@@ -351,15 +351,39 @@ The :mod:`test.test_support` module defines the following functions:
    .. versionadded:: 2.7
 
 
-.. function:: import_fresh_module(name, blocked_names=None, deprecated=False)
+.. function:: import_fresh_module(name, fresh=(), blocked=(), deprecated=False)
 
-   This function imports and returns a fresh copy of the named Python module. The
-   ``sys.modules`` cache is bypassed temporarily, and the ability to import the
-   modules named in *blocked_names* is suppressed for the duration of the import.
+   This function imports and returns a fresh copy of the named Python module
+   by removing the named module from ``sys.modules`` before doing the import.
+   Note that unlike :func:`reload`, the original module is not affected by
+   this operation.
+
+   *fresh* is an iterable of additional module names that are also removed
+   from the ``sys.modules`` cache before doing the import.
+
+   *blocked* is an iterable of module names that are replaced with :const:`0`
+   in the module cache during the import to ensure that attempts to import
+   them raise :exc:`ImportError`.
+
+   The named module and any modules named in the *fresh* and *blocked*
+   parameters are saved before starting the import and then reinserted into
+   ``sys.modules`` when the fresh import is complete.
 
    Module and package deprecation messages are suppressed during this import
    if *deprecated* is :const:`True`.
 
+   This function will raise :exc:`unittest.SkipTest` is the named module
+   cannot be imported.
+
+   Example use::
+
+      # Get copies of the warnings module for testing without
+      # affecting the version being used by the rest of the test suite
+      # One copy uses the C implementation, the other is forced to use
+      # the pure Python fallback implementation
+      py_warnings = import_fresh_module('warnings', blocked=['_warnings'])
+      c_warnings = import_fresh_module('warnings', fresh=['_warnings'])
+
    .. versionadded:: 2.7
 
 
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 625852dabc9..7e3c7a1acf5 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -8,7 +8,7 @@ import sys
 # We do a bit of trickery here to be able to test both the C implementation
 # and the Python implementation of the module.
 import heapq as c_heapq
-py_heapq = test_support.import_fresh_module('heapq', ['_heapq'])
+py_heapq = test_support.import_fresh_module('heapq', blocked=['_heapq'])
 
 class TestHeap(unittest.TestCase):
     module = None
diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
index f77df703715..5973e07c1b3 100644
--- a/Lib/test/test_support.py
+++ b/Lib/test/test_support.py
@@ -70,12 +70,43 @@ def import_module(name, deprecated=False):
             raise unittest.SkipTest(str(msg))
 
 
-def import_fresh_module(name, blocked_names=None, deprecated=False):
+def _save_and_remove_module(name, orig_modules):
+    """Helper function to save and remove a module from sys.modules
+
+       Return value is True if the module was in sys.modules and
+       False otherwise."""
+    saved = True
+    try:
+        orig_modules[name] = sys.modules[name]
+    except KeyError:
+        saved = False
+    else:
+        del sys.modules[name]
+    return saved
+
+
+def _save_and_block_module(name, orig_modules):
+    """Helper function to save and block a module in sys.modules
+
+       Return value is True if the module was in sys.modules and
+       False otherwise."""
+    saved = True
+    try:
+        orig_modules[name] = sys.modules[name]
+    except KeyError:
+        saved = False
+    sys.modules[name] = 0
+    return saved
+
+
+def import_fresh_module(name, fresh=(), blocked=(), deprecated=False):
     """Imports and returns a module, deliberately bypassing the sys.modules cache
     and importing a fresh copy of the module. Once the import is complete,
     the sys.modules cache is restored to its original state.
 
-    Importing of modules named in blocked_names is prevented while the fresh import
+    Modules named in fresh are also imported anew if needed by the import.
+
+    Importing of modules named in blocked is prevented while the fresh import
     takes place.
 
     If deprecated is True, any module or package deprecation messages
@@ -83,21 +114,24 @@ def import_fresh_module(name, blocked_names=None, deprecated=False):
     # NOTE: test_heapq and test_warnings include extra sanity checks to make
     # sure that this utility function is working as expected
     with _ignore_deprecated_imports(deprecated):
-        if blocked_names is None:
-            blocked_names = ()
+        # Keep track of modules saved for later restoration as well
+        # as those which just need a blocking entry removed
         orig_modules = {}
-        if name in sys.modules:
-            orig_modules[name] = sys.modules[name]
-            del sys.modules[name]
+        names_to_remove = []
+        _save_and_remove_module(name, orig_modules)
         try:
-            for blocked in blocked_names:
-                orig_modules[blocked] = sys.modules[blocked]
-                sys.modules[blocked] = 0
-            py_module = importlib.import_module(name)
+            for fresh_name in fresh:
+                _save_and_remove_module(fresh_name, orig_modules)
+            for blocked_name in blocked:
+                if not _save_and_block_module(blocked_name, orig_modules):
+                    names_to_remove.append(blocked_name)
+            fresh_module = importlib.import_module(name)
         finally:
-            for blocked, module in orig_modules.items():
-                sys.modules[blocked] = module
-        return py_module
+            for orig_name, module in orig_modules.items():
+                sys.modules[orig_name] = module
+            for name_to_remove in names_to_remove:
+                del sys.modules[name_to_remove]
+        return fresh_module
 
 
 def get_attribute(obj, name):
diff --git a/Lib/test/test_warnings.py b/Lib/test/test_warnings.py
index 4c3b43aec4b..9764783b53f 100644
--- a/Lib/test/test_warnings.py
+++ b/Lib/test/test_warnings.py
@@ -10,8 +10,8 @@ import warning_tests
 
 import warnings as original_warnings
 
-py_warnings = test_support.import_fresh_module('warnings', ['_warnings'])
-c_warnings = test_support.import_fresh_module('warnings')
+py_warnings = test_support.import_fresh_module('warnings', blocked=['_warnings'])
+c_warnings = test_support.import_fresh_module('warnings', fresh=['_warnings'])
 
 @contextmanager
 def warnings_state(module):
diff --git a/Misc/NEWS b/Misc/NEWS
index 74af1470974..aa4fa6d6a5a 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -839,6 +839,12 @@ Extension Modules
 Tests
 -----
 
+- Issue #5354: New test support function import_fresh_module() makes
+  it easy to import both normal and optimised versions of modules.
+  test_heapq and test_warnings have been adjusted to use it, tests for
+  other modules with both C and Python implementations in the stdlib
+  can be adjusted to use it over time.
+
 - Fix test_warnings to no longer reset the warnings filter.
 
 - Fix test_logging to no longer reset the warnings filter.