instrument_make.py

# Place in a Python package with this in the setup.py:
#
#     entry_points={
#         'console_scripts': [
#             'make = SOME_MODULE.instrument_make:wrapper',
#         ],
#     },


import re
import os
import subprocess
import sys


def get_non_venv_path():
    """
    Return the PATH environment variable with the VIRTUAL_ENV's bin path removed.
    """
    # There's no access to $_OLD_VIRTUAL_PATH (it's not exported) so
    # just manually remove the virtualenv's bin directory from the
    # PATH.
    venv_bin_re = re.escape(os.path.join(os.environ['VIRTUAL_ENV'], 'bin'))
    # Remove all instances from beginning, end, or middle -- but
    # always remove exactly one colon from whatever position it's
    # found in.  It should always be here just once and only at the
    # beginning, but be liberal here.  Using two scans here instead of
    # `venv:|:venv` regex otherwise there are some unlikely situations
    # where multiple-removal would fail, such as
    # `VENV:VENV:foo:bar:VENV:VENV` (would always leave one VENV no
    # matter which order the alternation list was in.)
    paths = os.environ['PATH']
    paths = re.sub(venv_bin_re + ':', '', paths)
    paths = re.sub(':' + venv_bin_re, '', paths)
    return paths


def find_non_venv_make():
    """
    Return the path to the make executable that isn't shadowed by
    the virtualenv.
    """
    outer_env = {}
    outer_env.update(os.environ)
    outer_env['PATH'] = get_non_venv_path()
    p = subprocess.run(["which make"], shell=True, capture_output=True, env=outer_env)
    return p.stdout.decode().strip()


def call_make(args):
    """
    Call make with given args and return exit code.
    """
    # Find real make, and call with virtualenv still active (because
    # it will be needed for most make targets.) However, add a special
    # marker environment variable so that we don't process recursive
    # calls to make.
    marked_env = {}
    marked_env.update(os.environ)
    marked_env['MAKE_INSTRUMENTED'] = 'true'
    p = subprocess.run([find_non_venv_make()] + list(args), env=marked_env)
    return p.returncode


def wrapper():
    """
    Entry point for CLI -- wraps a call to `make` and provides.

    Intended to be called with sys.argv bearing the arguments (via
    console_scripts in setup.py).
    """
    args = sys.argv[1:]

    if os.environ.get('MAKE_INSTRUMENTED') == 'true':
        print("This is a recursive call that should not be instrumented")
        returncode = call_make(args)
    else:
        print("This is an outer call that should be instrumented")
        # TODO: timing code, etc.
        returncode = call_make(args)
        # TODO: write metrics somewhere

    sys.exit(returncode)