Adjust distutils shim when removing _distutils_importer

1 files changed, 16 insertions, 3 deletions

diff --git a/setuptools/sandbox.py b/setuptools/sandbox.py
index 93ae8eb4..342a713f 100644
--- a/setuptools/sandbox.py
+++ b/setuptools/sandbox.py
@@ -185,8 +185,8 @@ def setup_context(setup_dir):
     temp_dir = os.path.join(setup_dir, 'temp')
     with save_pkg_resources_state():
         with save_modules():
-            hide_setuptools()
             with save_path():
+                hide_setuptools()
                 with save_argv():
                     with override_temp(temp_dir):
                         with pushd(setup_dir):
@@ -195,6 +195,15 @@ def setup_context(setup_dir):
                             yield
 
 
+_MODULES_TO_HIDE = {
+    'setuptools',
+    'distutils',
+    'pkg_resources',
+    'Cython',
+    '_distutils_importer',
+}
+
+
 def _needs_hiding(mod_name):
     """
     >>> _needs_hiding('setuptools')
@@ -212,8 +221,8 @@ def _needs_hiding(mod_name):
     >>> _needs_hiding('Cython')
     True
     """
-    pattern = re.compile(r'(setuptools|pkg_resources|distutils|Cython)(\.|$)')
-    return bool(pattern.match(mod_name))
+    base_module = mod_name.split('.', 1)[0]
+    return base_module in _MODULES_TO_HIDE
 
 
 def hide_setuptools():
@@ -223,6 +232,10 @@ def hide_setuptools():
     necessary to avoid issues such as #315 where setuptools upgrading itself
     would fail to find a function declared in the metadata.
     """
+    _distutils_importer = sys.modules.get('_distutils_importer', None)
+    if _distutils_importer is not None:
+        _distutils_importer.remove_shim()
+
     modules = filter(_needs_hiding, sys.modules)
     _clear_modules(modules)