gh-104310: Add importlib.util.allowing_all_extensions() (gh-104311)

(I'll be adding docs for this separately.)

1 files changed, 121 insertions, 0 deletions

diff --git a/Lib/test/test_importlib/test_util.py b/Lib/test/test_importlib/test_util.py
index 08a615e..0be5049 100644
--- a/Lib/test/test_importlib/test_util.py
+++ b/Lib/test/test_importlib/test_util.py
@@ -8,14 +8,29 @@ importlib_util = util.import_importlib('importlib.util')
 import importlib.util
 import os
 import pathlib
+import re
 import string
 import sys
 from test import support
+import textwrap
 import types
 import unittest
 import unittest.mock
 import warnings
 
+try:
+    import _testsinglephase
+except ImportError:
+    _testsinglephase = None
+try:
+    import _testmultiphase
+except ImportError:
+    _testmultiphase = None
+try:
+    import _xxsubinterpreters as _interpreters
+except ModuleNotFoundError:
+    _interpreters = None
+
 
 class DecodeSourceBytesTests:
 
@@ -637,5 +652,111 @@ class MagicNumberTests(unittest.TestCase):
         self.assertEqual(EXPECTED_MAGIC_NUMBER, actual, msg)
 
 
+@unittest.skipIf(_interpreters is None, 'subinterpreters required')
+class AllowingAllExtensionsTests(unittest.TestCase):
+
+    ERROR = re.compile("^<class 'ImportError'>: module (.*) does not support loading in subinterpreters")
+
+    def run_with_own_gil(self, script):
+        interpid = _interpreters.create(isolated=True)
+        try:
+            _interpreters.run_string(interpid, script)
+        except _interpreters.RunFailedError as exc:
+            if m := self.ERROR.match(str(exc)):
+                modname, = m.groups()
+                raise ImportError(modname)
+
+    def run_with_shared_gil(self, script):
+        interpid = _interpreters.create(isolated=False)
+        try:
+            _interpreters.run_string(interpid, script)
+        except _interpreters.RunFailedError as exc:
+            if m := self.ERROR.match(str(exc)):
+                modname, = m.groups()
+                raise ImportError(modname)
+
+    @unittest.skipIf(_testsinglephase is None, "test requires _testsinglephase module")
+    def test_single_phase_init_module(self):
+        script = textwrap.dedent('''
+            import importlib.util
+            with importlib.util.allowing_all_extensions():
+                import _testsinglephase
+            ''')
+        with self.subTest('check disabled, shared GIL'):
+            self.run_with_shared_gil(script)
+        with self.subTest('check disabled, per-interpreter GIL'):
+            self.run_with_own_gil(script)
+
+        script = textwrap.dedent(f'''
+            import importlib.util
+            with importlib.util.allowing_all_extensions(False):
+                import _testsinglephase
+            ''')
+        with self.subTest('check enabled, shared GIL'):
+            with self.assertRaises(ImportError):
+                self.run_with_shared_gil(script)
+        with self.subTest('check enabled, per-interpreter GIL'):
+            with self.assertRaises(ImportError):
+                self.run_with_own_gil(script)
+
+    @unittest.skipIf(_testmultiphase is None, "test requires _testmultiphase module")
+    def test_incomplete_multi_phase_init_module(self):
+        prescript = textwrap.dedent(f'''
+            from importlib.util import spec_from_loader, module_from_spec
+            from importlib.machinery import ExtensionFileLoader
+
+            name = '_test_shared_gil_only'
+            filename = {_testmultiphase.__file__!r}
+            loader = ExtensionFileLoader(name, filename)
+            spec = spec_from_loader(name, loader)
+
+            ''')
+
+        script = prescript + textwrap.dedent('''
+            import importlib.util
+            with importlib.util.allowing_all_extensions():
+                module = module_from_spec(spec)
+                loader.exec_module(module)
+            ''')
+        with self.subTest('check disabled, shared GIL'):
+            self.run_with_shared_gil(script)
+        with self.subTest('check disabled, per-interpreter GIL'):
+            self.run_with_own_gil(script)
+
+        script = prescript + textwrap.dedent('''
+            import importlib.util
+            with importlib.util.allowing_all_extensions(False):
+                module = module_from_spec(spec)
+                loader.exec_module(module)
+            ''')
+        with self.subTest('check enabled, shared GIL'):
+            self.run_with_shared_gil(script)
+        with self.subTest('check enabled, per-interpreter GIL'):
+            with self.assertRaises(ImportError):
+                self.run_with_own_gil(script)
+
+    @unittest.skipIf(_testmultiphase is None, "test requires _testmultiphase module")
+    def test_complete_multi_phase_init_module(self):
+        script = textwrap.dedent('''
+            import importlib.util
+            with importlib.util.allowing_all_extensions():
+                import _testmultiphase
+            ''')
+        with self.subTest('check disabled, shared GIL'):
+            self.run_with_shared_gil(script)
+        with self.subTest('check disabled, per-interpreter GIL'):
+            self.run_with_own_gil(script)
+
+        script = textwrap.dedent(f'''
+            import importlib.util
+            with importlib.util.allowing_all_extensions(False):
+                import _testmultiphase
+            ''')
+        with self.subTest('check enabled, shared GIL'):
+            self.run_with_shared_gil(script)
+        with self.subTest('check enabled, per-interpreter GIL'):
+            self.run_with_own_gil(script)
+
+
 if __name__ == '__main__':
     unittest.main()