bpo-30389 Adds detection of VS 2017 to distutils._msvccompiler (#1632)

2 files changed, 94 insertions, 29 deletions

diff --git a/Lib/distutils/_msvccompiler.py b/Lib/distutils/_msvccompiler.py
index b120273..ef1356b 100644
--- a/Lib/distutils/_msvccompiler.py
+++ b/Lib/distutils/_msvccompiler.py
@@ -17,6 +17,7 @@ import os
 import shutil
 import stat
 import subprocess
+import winreg
 
 from distutils.errors import DistutilsExecError, DistutilsPlatformError, \
                              CompileError, LibError, LinkError
@@ -24,10 +25,9 @@ from distutils.ccompiler import CCompiler, gen_lib_options
 from distutils import log
 from distutils.util import get_platform
 
-import winreg
 from itertools import count
 
-def _find_vcvarsall(plat_spec):
+def _find_vc2015():
     try:
         key = winreg.OpenKeyEx(
             winreg.HKEY_LOCAL_MACHINE,
@@ -38,9 +38,9 @@ def _find_vcvarsall(plat_spec):
         log.debug("Visual C++ is not registered")
         return None, None
 
+    best_version = 0
+    best_dir = None
     with key:
-        best_version = 0
-        best_dir = None
         for i in count():
             try:
                 v, vc_dir, vt = winreg.EnumValue(key, i)
@@ -53,25 +53,74 @@ def _find_vcvarsall(plat_spec):
                     continue
                 if version >= 14 and version > best_version:
                     best_version, best_dir = version, vc_dir
-        if not best_version:
-            log.debug("No suitable Visual C++ version found")
-            return None, None
+    return best_version, best_dir
+
+def _find_vc2017():
+    import _findvs
+    import threading
+
+    best_version = 0,   # tuple for full version comparisons
+    best_dir = None
+
+    # We need to call findall() on its own thread because it will
+    # initialize COM.
+    all_packages = []
+    def _getall():
+        all_packages.extend(_findvs.findall())
+    t = threading.Thread(target=_getall)
+    t.start()
+    t.join()
+
+    for name, version_str, path, packages in all_packages:
+        if 'Microsoft.VisualStudio.Component.VC.Tools.x86.x64' in packages:
+            vc_dir = os.path.join(path, 'VC', 'Auxiliary', 'Build')
+            if not os.path.isdir(vc_dir):
+                continue
+            try:
+                version = tuple(int(i) for i in version_str.split('.'))
+            except (ValueError, TypeError):
+                continue
+            if version > best_version:
+                best_version, best_dir = version, vc_dir
+    try:
+        best_version = best_version[0]
+    except IndexError:
+        best_version = None
+    return best_version, best_dir
 
-        vcvarsall = os.path.join(best_dir, "vcvarsall.bat")
-        if not os.path.isfile(vcvarsall):
-            log.debug("%s cannot be found", vcvarsall)
-            return None, None
+def _find_vcvarsall(plat_spec):
+    best_version, best_dir = _find_vc2017()
+    vcruntime = None
+    vcruntime_plat = 'x64' if 'amd64' in plat_spec else 'x86'
+    if best_version:
+        vcredist = os.path.join(best_dir, "..", "..", "redist", "MSVC", "**",
+            "Microsoft.VC141.CRT", "vcruntime140.dll")
+        try:
+            import glob
+            vcruntime = glob.glob(vcredist, recursive=True)[-1]
+        except (ImportError, OSError, LookupError):
+            vcruntime = None
+
+    if not best_version:
+        best_version, best_dir = _find_vc2015()
+        if best_version:
+            vcruntime = os.path.join(best_dir, 'redist', vcruntime_plat,
+                "Microsoft.VC140.CRT", "vcruntime140.dll")
+
+    if not best_version:
+        log.debug("No suitable Visual C++ version found")
+        return None, None
 
+    vcvarsall = os.path.join(best_dir, "vcvarsall.bat")
+    if not os.path.isfile(vcvarsall):
+        log.debug("%s cannot be found", vcvarsall)
+        return None, None
+
+    if not vcruntime or not os.path.isfile(vcruntime):
+        log.debug("%s cannot be found", vcruntime)
         vcruntime = None
-        vcruntime_spec = _VCVARS_PLAT_TO_VCRUNTIME_REDIST.get(plat_spec)
-        if vcruntime_spec:
-            vcruntime = os.path.join(best_dir,
-                vcruntime_spec.format(best_version))
-            if not os.path.isfile(vcruntime):
-                log.debug("%s cannot be found", vcruntime)
-                vcruntime = None
 
-        return vcvarsall, vcruntime
+    return vcvarsall, vcruntime
 
 def _get_vc_env(plat_spec):
     if os.getenv("DISTUTILS_USE_SDK"):
@@ -130,14 +179,6 @@ PLAT_TO_VCVARS = {
     'win-amd64' : 'x86_amd64',
 }
 
-# A map keyed by get_platform() return values to the file under
-# the VC install directory containing the vcruntime redistributable.
-_VCVARS_PLAT_TO_VCRUNTIME_REDIST = {
-    'x86' : 'redist\\x86\\Microsoft.VC{0}0.CRT\\vcruntime{0}0.dll',
-    'amd64' : 'redist\\x64\\Microsoft.VC{0}0.CRT\\vcruntime{0}0.dll',
-    'x86_amd64' : 'redist\\x64\\Microsoft.VC{0}0.CRT\\vcruntime{0}0.dll',
-}
-
 # A set containing the DLLs that are guaranteed to be available for
 # all micro versions of this Python version. Known extension
 # dependencies that are not in this set will be copied to the output
diff --git a/Lib/distutils/tests/test_msvccompiler.py b/Lib/distutils/tests/test_msvccompiler.py
index 4dc2488..70a9c93 100644
--- a/Lib/distutils/tests/test_msvccompiler.py
+++ b/Lib/distutils/tests/test_msvccompiler.py
@@ -76,12 +76,12 @@ class msvccompilerTestCase(support.TempdirManager,
         compiler = _msvccompiler.MSVCCompiler()
         compiler.initialize()
         dll = compiler._vcruntime_redist
-        self.assertTrue(os.path.isfile(dll))
+        self.assertTrue(os.path.isfile(dll), dll or "<None>")
 
         compiler._copy_vcruntime(tempdir)
 
         self.assertFalse(os.path.isfile(os.path.join(
-            tempdir, os.path.basename(dll))))
+            tempdir, os.path.basename(dll))), dll or "<None>")
 
     def test_get_vc_env_unicode(self):
         import distutils._msvccompiler as _msvccompiler
@@ -101,6 +101,30 @@ class msvccompilerTestCase(support.TempdirManager,
             if old_distutils_use_sdk:
                 os.environ['DISTUTILS_USE_SDK'] = old_distutils_use_sdk
 
+    def test_get_vc2017(self):
+        import distutils._msvccompiler as _msvccompiler
+
+        # This function cannot be mocked, so pass it if we find VS 2017
+        # and mark it skipped if we do not.
+        version, path = _msvccompiler._find_vc2017()
+        if version:
+            self.assertGreaterEqual(version, 15)
+            self.assertTrue(os.path.isdir(path))
+        else:
+            raise unittest.SkipTest("VS 2017 is not installed")
+
+    def test_get_vc2015(self):
+        import distutils._msvccompiler as _msvccompiler
+
+        # This function cannot be mocked, so pass it if we find VS 2015
+        # and mark it skipped if we do not.
+        version, path = _msvccompiler._find_vc2015()
+        if version:
+            self.assertGreaterEqual(version, 14)
+            self.assertTrue(os.path.isdir(path))
+        else:
+            raise unittest.SkipTest("VS 2015 is not installed")
+
 def test_suite():
     return unittest.makeSuite(msvccompilerTestCase)