monkeypatch_distutils.py

'''
Enhances distutils with:
    - better C++ support
    - incremental builds
    - parallel builds
'''

import os
import shlex
import subprocess
import sys
import distutils.sysconfig
import distutils.ccompiler
import distutils.unixccompiler
from distutils import log
from distutils.errors import (DistutilsExecError, DistutilsPlatformError,
                              CompileError, LibError, LinkError)
from multiprocessing.pool import ThreadPool

# threaded parallel map (optional)
parallel_jobs = 1

try:
    import _osx_support
    _osx_support._UNIVERSAL_CONFIG_VARS += ('CXXFLAGS', 'LDCXXSHARED',)
    _osx_support._COMPILER_CONFIG_VARS += ('LDCXXSHARED',)
except ImportError:
    _osx_support = None

distutils.unixccompiler.UnixCCompiler.executables.update({
    'compiler_cxx'      : ["c++"],
    'compiler_so_cxx'   : ["c++"],
    'linker_so_cxx'     : ["c++", "-shared"],
    'linker_exe_cxx'    : ["c++"],
})


def set_parallel_jobs(N):
    '''
    Set the number of parallel build jobs.
    N=1 : single threaded
    N=0 : use number of CPUs
    '''
    global parallel_jobs

    parallel_jobs = N or None


def incremental_parallel_compile(single_compile, objects, force):
    '''
    Call `single_compile` on each item in `objects` (in parallel) if
    any dependency (source file) is newer than the object file.
    '''
    mtimes = {}

    def deps(obj):
        # parse .d file (Makefile syntax)
        with open(os.path.splitext(obj)[0] + '.d') as handle:
            contents = handle.read().split(': ', 1)[-1]
            contents = contents.replace('\\\n', ' ')
            for dep in shlex.split(contents):
                yield dep

    def need_compile(obj):
        if force:
            return True

        try:
            obj_mtime = os.path.getmtime(obj)
            for dep in deps(obj):
                if dep not in mtimes:
                    mtimes[dep] = os.path.getmtime(dep)
                if obj_mtime < mtimes[dep]:
                    return True
        except EnvironmentError:
            return True

        return False

    with ThreadPool(parallel_jobs) as p:
        pmap = map if parallel_jobs == 1 else p.map
        for _ in pmap(single_compile, filter(need_compile, objects)):
            pass


def monkeypatch(parent, name):
    '''
    Decorator to replace a function or class method. Makes the
    unpatched function available as <patchedfunction>._super
    '''
    def wrapper(func):
        orig = getattr(parent, name)
        func._super = orig
        func.__name__ = name
        setattr(parent, name, func)
        return func
    return wrapper


def strip_broken_isysroot(args, start=0):
    '''
    Strip -isysroot which don't exist from args.

    Those can come from lib/python2.7/_sysconfigdata_x86_64_apple_darwin13_4_0.py
    '''
    assert isinstance(args, list)

    try:
        i = args.index('-isysroot', start)
    except ValueError:
        return

    strip_broken_isysroot(args, i + 2)

    if not os.path.isdir(args[i + 1]):
        args[i:i + 2] = []


@monkeypatch(distutils.sysconfig, 'customize_compiler')
def customize_compiler(compiler):
    # remove problematic flags
    if sys.platform == 'linux' and (
            'icpc' in os.getenv('CXX', '') or
            'clang' in os.getenv('CC', '') or
            'clang' in os.getenv('LD', '')):
        import re
        re_flto = re.compile(r'-flto\S*|-fno-semantic-interposition')
        config_vars = distutils.sysconfig.get_config_vars()
        for (key, value) in config_vars.items():
            if re_flto.search(str(value)) is not None:
                config_vars[key] = re_flto.sub('', value)

    customize_compiler._super(compiler)

    if compiler.compiler_type != "unix":
        return

    (cxx, ccshared, ldcxxshared) = \
            distutils.sysconfig.get_config_vars('CXX', 'CCSHARED', 'LDCXXSHARED')

    cxx = os.environ.get('CXX') or cxx
    cxxflags = os.environ.get('CXXFLAGS', '') + ' ' + os.environ.get('CPPFLAGS', '')
    ldcxxshared = os.environ.get('LDCXXSHARED', ldcxxshared or '') + \
            ' ' + os.environ.get('LDFLAGS', '') + \
            ' ' + os.environ.get('CXXFLAGS', '') + \
            ' ' + os.environ.get('CPPFLAGS', '')

    cxx_cmd = cxx + ' ' + cxxflags

    # C++17 by default
    if '-std=' not in cxx_cmd:
        cxx_cmd += ' -std=c++17'

    compiler.set_executables(
            compiler_cxx=cxx_cmd,
            compiler_so_cxx=cxx_cmd + ' ' + ccshared,
            linker_so_cxx=ldcxxshared,
            linker_exe_cxx=cxx)

@monkeypatch(distutils.unixccompiler.UnixCCompiler, 'compile')
def compile(self, sources, output_dir=None, macros=None,
        include_dirs=None, debug=0, extra_preargs=None, extra_postargs=None,
        depends=None):
    '''
    Enable parallel and incremental build.

    To do a clean build, please remove the "build" directory.
    '''
    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
            output_dir, macros, include_dirs, sources, depends, extra_postargs)
    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)

    compiler_so = self.compiler_so
    compiler_so_cxx = self.compiler_so_cxx

    # strips non-existing -isysroot
    strip_broken_isysroot(compiler_so)
    strip_broken_isysroot(compiler_so_cxx)

    if sys.platform == 'darwin' and _osx_support is not None:
        # strips duplicated -isysroot
        compiler_so = _osx_support.compiler_fixup(compiler_so, cc_args + extra_postargs)
        compiler_so_cxx = _osx_support.compiler_fixup(compiler_so_cxx, cc_args + extra_postargs)

    # generate dependency (.d) file
    cc_args.append('-MMD')

    def _single_compile(obj):
        src = build[obj][0]

        # _compile
        compiler = compiler_so_cxx \
                if self.detect_language(src) == 'c++' \
                else compiler_so
        try:
            self.spawn(compiler + cc_args + [src, '-o', obj] + extra_postargs)
        except distutils.errors.DistutilsExecError as msg:
            raise distutils.errors.CompileError(msg)

    incremental_parallel_compile(_single_compile, objects, self.force)

    return objects

if sys.platform.startswith('win'):
    try:
        from distutils import _msvccompiler as msvccompiler
    except ImportError:
        from distutils import msvccompiler

    @monkeypatch(msvccompiler.MSVCCompiler, 'compile')
    def compile(self, sources,
                output_dir=None, macros=None, include_dirs=None, debug=0,
                extra_preargs=None, extra_postargs=None, depends=None):
        '''
        Enable parallel and incremental build.
        '''
        if not self.initialized:
            self.initialize()
        compile_info = self._setup_compile(output_dir, macros, include_dirs,
                                           sources, depends, extra_postargs)
        macros, objects, extra_postargs, pp_opts, build = compile_info

        compile_opts = extra_preargs or []
        compile_opts.append ('/c')
        if debug:
            compile_opts.extend(self.compile_options_debug)
        else:
            compile_opts.extend(self.compile_options)

        env = dict(os.environ)
        try:
            env['PATH'] = self._paths
        except AttributeError:
            pass

        def _single_compile(obj):
            try:
                src, ext = build[obj]
            except KeyError:
                return

            add_cpp_opts = False

            if debug:
                # pass the full pathname to MSVC in debug mode,
                # this allows the debugger to find the source file
                # without asking the user to browse for it
                src = os.path.abspath(src)

            if ext in self._c_extensions:
                input_opt = "/Tc" + src
            elif ext in self._cpp_extensions:
                input_opt = "/Tp" + src
                add_cpp_opts = True
            elif ext in self._rc_extensions:
                # compile .RC to .RES file
                input_opt = src
                output_opt = "/fo" + obj
                try:
                    self.spawn([self.rc] + pp_opts +
                               [output_opt] + [input_opt])
                except DistutilsExecError as msg:
                    raise CompileError(msg)
                return
            elif ext in self._mc_extensions:
                # Compile .MC to .RC file to .RES file.
                h_dir = os.path.dirname(src)
                rc_dir = os.path.dirname(obj)
                try:
                    # first compile .MC to .RC and .H file
                    self.spawn([self.mc] +
                               ['-h', h_dir, '-r', rc_dir] + [src])
                    base, _ = os.path.splitext (os.path.basename (src))
                    rc_file = os.path.join (rc_dir, base + '.rc')
                    # then compile .RC to .RES file
                    self.spawn([self.rc] +
                               ["/fo" + obj] + [rc_file])

                except DistutilsExecError as msg:
                    raise CompileError(msg)
                return
            else:
                # how to handle this file?
                raise CompileError("Don't know how to compile %s to %s"
                                   % (src, obj))

            args = [self.cc] + compile_opts + pp_opts
            if add_cpp_opts:
                args.append('/EHsc')
            args.append(input_opt)
            args.append("/Fo" + obj)
            args.extend(extra_postargs)

            try:
                msvc_spawn_and_write_d_file(obj, src, args, env, self.dry_run)
            except DistutilsExecError as msg:
                raise CompileError(msg)

        incremental_parallel_compile(_single_compile, objects, self.force)

        return objects

    def msvc_spawn_and_write_d_file(obj, src, cmd, env, dry_run):
        '''
        Run command with /showIncludes and convert output to dependency (.d) file.
        '''
        log.info(' '.join(cmd))

        if dry_run:
            return

        process = subprocess.Popen(cmd + ['/showIncludes'], env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                universal_newlines=True)
        out = process.communicate()[0]

        deps = set([src])

        for line in out.splitlines():
            if not line.startswith('Note: including file:'):
                sys.stderr.write(line + '\n')
                continue

            dep = line[21:].strip()
            dep_lower = dep.lower()
            if not (
                    # filter out system headers
                    'microsoft visual studio' in dep_lower or
                    'windows kits' in dep_lower):
                deps.add(dep)

        if process.returncode != 0:
            raise DistutilsExecError("command %r failed with exit status %d"
                    % (cmd, process.returncode))

        with open(os.path.splitext(obj)[0] + '.d', 'w') as handle:
            handle.write(': ')
            for dep in deps:
                handle.write(' \\\n')
                handle.write(dep.replace('\\', '\\\\').replace(' ', '\\ '))