Consolidate function patching and resolution of unpatched function, aligning pattern with the patched classes.

3 files changed, 32 insertions, 18 deletions

diff --git a/setuptools/monkey.py b/setuptools/monkey.py
index 33083831..24739d97 100644
--- a/setuptools/monkey.py
+++ b/setuptools/monkey.py
@@ -6,6 +6,8 @@ import sys
 import distutils.filelist
 import platform
 
+from .py26compat import import_module
+
 import setuptools
 
 
@@ -103,7 +105,20 @@ def _patch_distribution_metadata_write_pkg_info():
     )
 
 
-unpatched = dict()
+def patch_func(replacement, original):
+    # first set the 'unpatched' attribute on the replacement to
+    # point to the original.
+    vars(replacement).setdefault('unpatched', original)
+
+    # next resolve the module in which the original func resides
+    target_mod = import_module(original.__module__)
+
+    # finally replace the function in the original module
+    setattr(target_mod, original.__name__, replacement)
+
+
+def get_unpatched_func(candidate):
+    return getattr(candidate, 'unpatched')
 
 
 def patch_for_msvc_specialized_compiler():
@@ -129,29 +144,21 @@ def patch_for_msvc_specialized_compiler():
         # Compilers only availables on Microsoft Windows
         return
 
-    if unpatched:
-        # Already patched
-        return
-
     try:
         # Patch distutils.msvc9compiler
-        unpatched['msvc9_find_vcvarsall'] = msvc9compiler.find_vcvarsall
-        msvc9compiler.find_vcvarsall = msvc.msvc9_find_vcvarsall
-        unpatched['msvc9_query_vcvarsall'] = msvc9compiler.query_vcvarsall
-        msvc9compiler.query_vcvarsall = msvc.msvc9_query_vcvarsall
+        patch_func(msvc.msvc9_find_vcvarsall, msvc9compiler.find_vcvarsall)
+        patch_func(msvc.msvc9_query_vcvarsall, msvc9compiler.query_vcvarsall)
     except NameError:
         pass
 
     try:
         # Patch distutils._msvccompiler._get_vc_env
-        unpatched['msvc14_get_vc_env'] = msvc14compiler._get_vc_env
-        msvc14compiler._get_vc_env = msvc.msvc14_get_vc_env
+        patch_func(msvc.msvc14_get_vc_env, msvc14compiler._get_vc_env)
     except NameError:
         pass
 
     try:
         # Patch distutils._msvccompiler.gen_lib_options for Numpy
-        unpatched['msvc14_gen_lib_options'] = msvc14compiler.gen_lib_options
-        msvc14compiler.gen_lib_options = msvc.msvc14_gen_lib_options
+        patch_func(msvc.msvc14_gen_lib_options, msvc14compiler.gen_lib_options)
     except NameError:
         pass
diff --git a/setuptools/msvc.py b/setuptools/msvc.py
index a902a4a7..ae5a2b6a 100644
--- a/setuptools/msvc.py
+++ b/setuptools/msvc.py
@@ -24,7 +24,7 @@ from pkg_resources.extern.packaging.version import LegacyVersion
 
 from setuptools.extern.six.moves import filterfalse
 
-from . import monkey
+from .monkey import get_unpatched_func
 
 if platform.system() == 'Windows':
     from setuptools.extern.six.moves import winreg
@@ -87,7 +87,7 @@ def msvc9_find_vcvarsall(version):
         if os.path.isfile(vcvarsall):
             return vcvarsall
 
-    return monkey.unpatched['msvc9_find_vcvarsall'](version)
+    return get_unpatched_func(msvc9_find_vcvarsall)(version)
 
 
 def msvc9_query_vcvarsall(ver, arch='x86', *args, **kwargs):
@@ -120,7 +120,7 @@ def msvc9_query_vcvarsall(ver, arch='x86', *args, **kwargs):
     """
     # Try to get environement from vcvarsall.bat (Classical way)
     try:
-        orig = monkey.unpatched['msvc9_query_vcvarsall']
+        orig = get_unpatched_func(msvc9_query_vcvarsall)
         return orig(ver, arch, *args, **kwargs)
     except distutils.errors.DistutilsPlatformError:
         # Pass error if Vcvarsall.bat is missing
@@ -160,7 +160,7 @@ def msvc14_get_vc_env(plat_spec):
     """
     # Try to get environment from vcvarsall.bat (Classical way)
     try:
-        return monkey.unpatched['msvc14_get_vc_env'](plat_spec)
+        return get_unpatched_func(msvc14_get_vc_env)(plat_spec)
     except distutils.errors.DistutilsPlatformError:
         # Pass error Vcvarsall.bat is missing
         pass
@@ -183,7 +183,7 @@ def msvc14_gen_lib_options(*args, **kwargs):
         import numpy as np
         if LegacyVersion(np.__version__) < LegacyVersion('1.11.2'):
             return np.distutils.ccompiler.gen_lib_options(*args, **kwargs)
-    return monkey.unpatched['msvc14_gen_lib_options'](*args, **kwargs)
+    return get_unpatched_func(msvc14_gen_lib_options)(*args, **kwargs)
 
 
 def _augment_exception(exc, version, arch=''):
diff --git a/setuptools/py26compat.py b/setuptools/py26compat.py
index 90cd695a..5778cdf1 100644
--- a/setuptools/py26compat.py
+++ b/setuptools/py26compat.py
@@ -22,3 +22,10 @@ def strip_fragment(url):
 
 if sys.version_info >= (2, 7):
     strip_fragment = lambda x: x
+
+
+try:
+    from importlib import import_module
+except ImportError:
+    def import_module(module_name):
+        return __import__(module_name, fromlist=['__name__'])