In distutils_hack, only add the metadata finder once. In ensure_local_distutils, rely on a context manager for reliable manipulation. Fixes #2958.

1 files changed, 16 insertions, 3 deletions

diff --git a/_distutils_hack/__init__.py b/_distutils_hack/__init__.py
index 22bc9ed6..85a51370 100644
--- a/_distutils_hack/__init__.py
+++ b/_distutils_hack/__init__.py
@@ -3,6 +3,7 @@ import os
 import re
 import importlib
 import warnings
+import contextlib
 
 
 is_pypy = '__pypy__' in sys.builtin_module_names
@@ -52,9 +53,8 @@ def ensure_local_distutils():
     # With the DistutilsMetaFinder in place,
     # perform an import to cause distutils to be
     # loaded from setuptools._distutils. Ref #2906.
-    add_shim()
-    importlib.import_module('distutils')
-    remove_shim()
+    with shim():
+        importlib.import_module('distutils')
 
     # check that submodules load as expected
     core = importlib.import_module('distutils.core')
@@ -129,6 +129,19 @@ class DistutilsMetaFinder:
 DISTUTILS_FINDER = DistutilsMetaFinder()
 
 
+def ensure_shim():
+    DISTUTILS_FINDER in sys.meta_path or add_shim()
+
+
+@contextlib.contextmanager
+def shim():
+    add_shim()
+    try:
+        yield
+    finally:
+        remove_shim()
+
+
 def add_shim():
     sys.meta_path.insert(0, DISTUTILS_FINDER)