msvc_env.py

#!/usr/bin/env python2
import argparse
import json
import logging
import nt
import os
import re
import subprocess
import sys
from distutils.spawn import find_executable

import six

try:
    from typing import Mapping, NoReturn, Optional
except ImportError:
    pass

VSWHERE_DEFAULT_PATH = (
    r"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe"
)
MSVC_COMPONENT_ID = "Microsoft.VisualStudio.Component.VC.Tools.x86.x64"


def find_visual_studio(vs_version, vswhere=VSWHERE_DEFAULT_PATH):
    # type: (str, str) -> Mapping[str, str]
    vswhere = os.path.expandvars(vswhere)
    version_filter = "[{},{})".format(vs_version, float(vs_version) + 1)
    vswhere_cmd = [
        vswhere,
        "-products",
        "*",
        "-format",
        "json",
        "-utf8",
        "-sort",
        "-requires",
        MSVC_COMPONENT_ID,
        "-version",
        version_filter,
    ]

    logging.debug("Running: %s", subprocess.list2cmdline(vswhere_cmd))
    results = json.loads(subprocess.check_output(vswhere_cmd))

    if not results:
        raise RuntimeError("Visual Studio was not found")

    return results[0]


def get_msvc_env(vs_info, arch, msvc_version=None, winsdk_version=None):
    # type: (Mapping[str, str], str, Optional[str], Optional[str]) -> Mapping[str, str]
    if arch == "x86":
        msvc_arch = "x64_x86"
    elif arch == "x86_64":
        msvc_arch = "x64"
    else:
        raise ValueError("Unknown arch: {}".format(arch))

    install_root = vs_info["installationPath"]
    env_script = os.path.join(install_root, r"VC\Auxiliary\Build\vcvarsall.bat")

    # Preserve case of env keys (os.environ unconditionally uppercases them on Windows).
    canonical_keys = {k.upper(): k for k in nt.environ}

    env = {canonical_keys.get(k.upper(), k): v for k, v in six.iteritems(os.environ)}
    env["VSCMD_SKIP_SENDTELEMETRY"] = "1"  # Don't send telemetry data every time

    vcvarsall_cmd = [env_script, msvc_arch]
    if winsdk_version:
        vcvarsall_cmd.append(winsdk_version)
    if msvc_version:
        vcvarsall_cmd.append("-vcvars_ver={}".format(msvc_version))

    logging.debug("Running: %s", subprocess.list2cmdline(vcvarsall_cmd))
    vcvarsall_output = subprocess.check_output(
        "{} && cl 2>nul && set".format(subprocess.list2cmdline(vcvarsall_cmd)),
        shell=True,
        # Disable ctrl-c otherwise cmd can hang instead of terminating cleanly.
        creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
        env=env,
    )

    vcvarsall_env = {}

    for line in vcvarsall_output.splitlines():
        m = re.match(r"([^=]+)=(.+)", six.ensure_str(line))
        if m is None:
            continue

        k, v = m.group(1), m.group(2)
        if os.environ.get(k) == v:
            continue
        if k.startswith("__VSCMD_") or k.startswith("VSCMD_"):
            continue

        if k.upper() == "PATH" and "PATH" in os.environ:
            # Remove existing paths
            v = re.sub(";?{};?".format(re.escape(os.environ["PATH"])), "", v)

        vcvarsall_env[k] = v

    return vcvarsall_env


def _make_env(vs_env):
    # type: (Mapping[str, str] -> Mapping[str, str])

    # Preserve case of env keys (os.environ unconditionally uppercases them on Windows).
    canonical_keys = {k.upper(): k for k in nt.environ}

    env = {canonical_keys.get(k.upper(), k): v for k, v in six.iteritems(os.environ)}

    for k, v in six.iteritems(vs_env):
        k = canonical_keys.get(k.upper(), k)
        if k.upper() == "PATH" and k in env:
            # Special case: prepend PATH.
            env[k] = "{};{}".format(v, env[k])
        else:
            env[k] = v

    return env


def main():
    # type: () -> NoReturn
    logging.basicConfig(level=logging.DEBUG)

    argv = sys.argv[1:]
    try:
        i = argv.index("--")
    except ValueError:
        cmd = []
    else:
        argv, cmd = argv[:i], argv[i + 1 :]

    parser = argparse.ArgumentParser()
    parser.add_argument("--vs_version", default="16")
    parser.add_argument("--arch", choices=["x86", "x86_64"], default="x86")
    parser.add_argument("--msvc_version")
    parser.add_argument("--winsdk_version")
    parser.add_argument("--vswhere", default=VSWHERE_DEFAULT_PATH)
    parser.add_argument("--json", action="store_true")
    parser.add_argument("cmd", nargs=argparse.REMAINDER)

    args = parser.parse_args(argv)
    args.cmd += cmd

    vs_info = find_visual_studio(vs_version=args.vs_version, vswhere=args.vswhere)
    msvc_env = get_msvc_env(
        vs_info,
        arch=args.arch,
        msvc_version=args.msvc_version,
        winsdk_version=args.winsdk_version,
    )

    if args.json:
        key = dict(
            vs_version=args.vs_version,
            arch=args.arch,
            msvc_version=args.msvc_version,
            winsdk_version=args.winsdk_version,
        )
        key_str = ",".join(
            "{}={}".format(k, v) for k, v in sorted(six.iteritems(key)) if v
        )
        json_doc = {key_str: {"product_info": vs_info, "env": msvc_env}}
        json.dump(json_doc, sys.stdout, sort_keys=True, indent=4)
        sys.exit(0)

    else:
        if not args.cmd:
            args.cmd = ["cmd", "/C", "set"]

        subproc_env = _make_env(msvc_env)

        # Resolve cmd[0] using PATH.
        path = next((v for k, v in six.iteritems(subproc_env) if k.upper() == "PATH"), "")
        args.cmd[0] = find_executable(args.cmd[0], path) or args.cmd[0]

        logging.debug("Running: %s", subprocess.list2cmdline(args.cmd))
        sys.exit(subprocess.call(args.cmd, env=subproc_env))


if __name__ == "__main__":
    main()