gh-152659: Address review feedback from brettcannon

Author: zainnadeem786

Lib/test/support/import_helper.py

@@ -105,6 +105,39 @@ def _save_and_remove_modules(names):
     return orig_modules
 
 
+_PARENT_ATTR_MISSING = object()
+
+
+def _save_parent_attrs(names):
+    parent_attrs = {}
+    prefixes = tuple(name + '.' for name in names)
+    removed = {modname for modname in sys.modules
+               if modname in names or modname.startswith(prefixes)}
+    for modname in set(names) | removed:
+        parent_name, _, attr = modname.rpartition('.')
+        if not parent_name:
+            continue
+        if parent_name in removed or parent_name.startswith(prefixes):
+            continue
+        if (parent := sys.modules.get(parent_name)) is None:
+            continue
+        parent_attrs[parent_name, attr] = getattr(parent, attr, _PARENT_ATTR_MISSING)
+    return parent_attrs
+
+
+def _restore_parent_attrs(parent_attrs):
+    for (parent_name, attr), value in parent_attrs.items():
+        if (parent := sys.modules.get(parent_name)) is None:
+            continue
+        if value is _PARENT_ATTR_MISSING:
+            try:
+                delattr(parent, attr)
+            except AttributeError:
+                pass
+        else:
+            setattr(parent, attr, value)
+
+
 @contextlib.contextmanager
 def frozen_modules(enabled=True):
     """Force frozen modules to be used (or not).
@@ -179,6 +212,7 @@ def import_fresh_module(name, fresh=(), blocked=(), *,
         fresh = list(fresh)
         blocked = list(blocked)
         names = {name, *fresh, *blocked}
+        orig_parent_attrs = _save_parent_attrs(names)
         orig_modules = _save_and_remove_modules(names)
         for modname in blocked:
             sys.modules[modname] = None
@@ -195,6 +229,7 @@ def import_fresh_module(name, fresh=(), blocked=(), *,
         finally:
             _save_and_remove_modules(names)
             sys.modules.update(orig_modules)
+            _restore_parent_attrs(orig_parent_attrs)
 
 
 class CleanImport(object):

Lib/test/test_support.py

@@ -91,6 +91,35 @@ def test_import_module(self):
     def test_import_fresh_module(self):
         import_helper.import_fresh_module("ftplib")
 
+    def test_import_fresh_module_restores_parent_attr(self):
+        import importlib.util
+
+        name = importlib.util.__name__
+        original_module = sys.modules[name]
+        self.assertIs(importlib.util, original_module)
+
+        fresh_module = import_helper.import_fresh_module(name)
+
+        self.assertIsNot(fresh_module, original_module)
+        self.assertIs(sys.modules[name], original_module)
+        self.assertIs(importlib.util, original_module)
+
+    def test_import_fresh_module_removes_added_parent_attr(self):
+        import xml
+
+        name = "xml.sax"
+        self.enterContext(import_helper.CleanImport(name))
+        if hasattr(xml, "sax"):
+            self.addCleanup(setattr, xml, "sax", xml.sax)
+            del xml.sax
+        self.assertFalse(hasattr(xml, "sax"))
+
+        fresh_module = import_helper.import_fresh_module(name)
+
+        self.assertIsNotNone(fresh_module)
+        self.assertNotIn(name, sys.modules)
+        self.assertFalse(hasattr(xml, "sax"))
+
     def test_get_attribute(self):
         self.assertEqual(support.get_attribute(self, "test_get_attribute"),
                         self.test_get_attribute)