Merge pull request RustPython#3488 from deantvv/frozen-import-helper

import_helper: refactor `frozen_modules`

Author: coolreader18

Lib/test/support/import_helper.py

@@ -91,204 +91,105 @@ def _save_and_remove_modules(names):
     return orig_modules
 
 
-# XXX RUSTPYTHON: need _imp._override_frozen_modules_for_tests
-# @contextlib.contextmanager
-# def frozen_modules(enabled=True):
-#     """Force frozen modules to be used (or not).
-
-#     This only applies to modules that haven't been imported yet.
-#     Also, some essential modules will always be imported frozen.
-#     """
-#     _imp._override_frozen_modules_for_tests(1 if enabled else -1)
-#     try:
-#         yield
-#     finally:
-#         _imp._override_frozen_modules_for_tests(0)
-
-
-# XXX RUSTPYTHON: new implementation needs fronzen_modules
-# def import_fresh_module(name, fresh=(), blocked=(), *,
-#                         deprecated=False,
-#                         usefrozen=False,
-#                         ):
-#     """Import and return a module, deliberately bypassing sys.modules.
-
-#     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 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. If one of these
-#     modules can't be imported, None is returned.
-
-#     *blocked* is an iterable of module names that are replaced with None
-#     in the module cache during the import to ensure that attempts to import
-#     them raise 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 True.
-
-#     This function will raise ImportError if the named module cannot be
-#     imported.
-
-#     If "usefrozen" is False (the default) then the frozen importer is
-#     disabled (except for essential modules like importlib._bootstrap).
-#     """
-#     # NOTE: test_heapq, test_json and test_warnings include extra sanity checks
-#     # to make sure that this utility function is working as expected
-#     with _ignore_deprecated_imports(deprecated):
-#         # Keep track of modules saved for later restoration as well
-#         # as those which just need a blocking entry removed
-#         fresh = list(fresh)
-#         blocked = list(blocked)
-#         names = {name, *fresh, *blocked}
-#         orig_modules = _save_and_remove_modules(names)
-#         for modname in blocked:
-#             sys.modules[modname] = None
-
-#         try:
-#             with frozen_modules(usefrozen):
-#                 # Return None when one of the "fresh" modules can not be imported.
-#                 try:
-#                     for modname in fresh:
-#                         __import__(modname)
-#                 except ImportError:
-#                     return None
-#                 return importlib.import_module(name)
-#         finally:
-#             _save_and_remove_modules(names)
-#             sys.modules.update(orig_modules)
-
-
-# TODO RUSTPYTHON: old implementation
-def _save_and_remove_module(name, orig_modules):
-    """Helper function to save and remove a module from sys.modules
-    Raise ImportError if the module can't be imported.
+# TODO RUSTPYTHON: need _imp._override_frozen_modules_for_tests
+# The following implementation is NOT correct and only raise
+# exception when it needs enabled=True
+@contextlib.contextmanager
+def frozen_modules(enabled=True):
+    """Force frozen modules to be used (or not).
+
+    This only applies to modules that haven't been imported yet.
+    Also, some essential modules will always be imported frozen.
     """
-    # try to import the module and raise an error if it can't be imported
-    if name not in sys.modules:
-        __import__(name)
-        del sys.modules[name]
-    for modname in list(sys.modules):
-        if modname == name or modname.startswith(name + '.'):
-            orig_modules[modname] = sys.modules[modname]
-            del sys.modules[modname]
+    if enabled:
+        raise NotImplemented("frozen_modules is not implemented on RustPython")
 
+    yield
 
-# TODO RUSTPYTHON: old implementation
-def _save_and_block_module(name, orig_modules):
-    """Helper function to save and block a module in sys.modules
-
-    Return True if the module was in sys.modules, False otherwise.
-    """
-    saved = True
-    try:
-        orig_modules[name] = sys.modules[name]
-    except KeyError:
-        saved = False
-    sys.modules[name] = None
-    return saved
+    # TODO: original implementation
+    # _imp._override_frozen_modules_for_tests(1 if enabled else -1)
+    # try:
+    #     yield
+    # finally:
+    #     _imp._override_frozen_modules_for_tests(0)
 
 
-# TODO RUSTPYTHON: old implementation
-def import_fresh_module(name, fresh=(), blocked=(), deprecated=False):
+# TODO: `frozen_modules` is not supported
+def import_fresh_module(name, fresh=(), blocked=(), *,
+                        deprecated=False,
+                        usefrozen=False,
+                        ):
     """Import and return a module, deliberately bypassing sys.modules.
+
     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 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.
+    from the sys.modules cache before doing the import. If one of these
+    modules can't be imported, None is returned.
+
     *blocked* is an iterable of module names that are replaced with None
     in the module cache during the import to ensure that attempts to import
     them raise 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 True.
+
     This function will raise ImportError if the named module cannot be
     imported.
+
+    If "usefrozen" is False (the default) then the frozen importer is
+    disabled (except for essential modules like importlib._bootstrap).
     """
     # NOTE: test_heapq, test_json and test_warnings include extra sanity checks
     # to make sure that this utility function is working as expected
     with _ignore_deprecated_imports(deprecated):
         # Keep track of modules saved for later restoration as well
         # as those which just need a blocking entry removed
-        orig_modules = {}
-        names_to_remove = []
-        _save_and_remove_module(name, orig_modules)
+        fresh = list(fresh)
+        blocked = list(blocked)
+        names = {name, *fresh, *blocked}
+        orig_modules = _save_and_remove_modules(names)
+        for modname in blocked:
+            sys.modules[modname] = None
+
         try:
-            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)
-        except ImportError:
-            fresh_module = None
+            with frozen_modules(usefrozen):
+                # Return None when one of the "fresh" modules can not be imported.
+                try:
+                    for modname in fresh:
+                        __import__(modname)
+                except ImportError:
+                    return None
+                return importlib.import_module(name)
         finally:
-            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
-
-# TODO RUSTPYTHON: new implementation needs fronzen_modules
-# class CleanImport(object):
-#     """Context manager to force import to return a new module reference.
-
-#     This is useful for testing module-level behaviours, such as
-#     the emission of a DeprecationWarning on import.
-
-#     Use like this:
-
-#         with CleanImport("foo"):
-#             importlib.import_module("foo") # new reference
-
-#     If "usefrozen" is False (the default) then the frozen importer is
-#     disabled (except for essential modules like importlib._bootstrap).
-#     """
-
-#     def __init__(self, *module_names, usefrozen=False):
-#         self.original_modules = sys.modules.copy()
-#         for module_name in module_names:
-#             if module_name in sys.modules:
-#                 module = sys.modules[module_name]
-#                 # It is possible that module_name is just an alias for
-#                 # another module (e.g. stub for modules renamed in 3.x).
-#                 # In that case, we also need delete the real module to clear
-#                 # the import cache.
-#                 if module.__name__ != module_name:
-#                     del sys.modules[module.__name__]
-#                 del sys.modules[module_name]
-#         self._frozen_modules = frozen_modules(usefrozen)
-
-#     def __enter__(self):
-#         self._frozen_modules.__enter__()
-#         return self
-
-#     def __exit__(self, *ignore_exc):
-#         sys.modules.update(self.original_modules)
-#         self._frozen_modules.__exit__(*ignore_exc)
-
-
-# TODO RUSTPYTHON: old implementation
+            _save_and_remove_modules(names)
+            sys.modules.update(orig_modules)
+
+
+# TODO: `frozen_modules` is not supported
 class CleanImport(object):
     """Context manager to force import to return a new module reference.
+
     This is useful for testing module-level behaviours, such as
     the emission of a DeprecationWarning on import.
+
     Use like this:
+
         with CleanImport("foo"):
             importlib.import_module("foo") # new reference
+
+    If "usefrozen" is False (the default) then the frozen importer is
+    disabled (except for essential modules like importlib._bootstrap).
     """
 
-    def __init__(self, *module_names):
+    def __init__(self, *module_names, usefrozen=False):
         self.original_modules = sys.modules.copy()
         for module_name in module_names:
             if module_name in sys.modules:
@@ -300,12 +201,15 @@ def __init__(self, *module_names):
                 if module.__name__ != module_name:
                     del sys.modules[module.__name__]
                 del sys.modules[module_name]
+        self._frozen_modules = frozen_modules(usefrozen)
 
     def __enter__(self):
+        self._frozen_modules.__enter__()
         return self
 
     def __exit__(self, *ignore_exc):
         sys.modules.update(self.original_modules)
+        self._frozen_modules.__exit__(*ignore_exc)
 
 
 class DirsOnSysPath(object):