configure

#!/usr/bin/env python

import argparse
import pathlib
import platform
import shutil
import subprocess
import sys

BAZELRC = '.bazelrc.configure'

UNIX_PATH = 'PATH=/usr/local/miniconda3/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'


def is_linux():
    return platform.system() == 'Linux'


def is_macos():
    return platform.system() == 'Darwin'


def is_windows():
    return platform.system() == 'Windows'


def write_to_bazelrc(line):
    with open(BAZELRC, 'a') as f:
        f.write(line + '\n')


def choice_prompt(question, choices, default):
    inp = ""
    while not inp in choices:
        inp = input("{0}? ({1})[{2}]: ".format(question, ",".join(choices), default))
        if not inp:
            inp = default
        elif inp not in choices:
            print("Invalid choice: {}".format(inp))
    return inp


class Configure:

    def __init__(self):
        conda = shutil.which('conda')
        if not conda:
            print('Please install conda.')
            print('See: https://docs.conda.io/projects/conda/en/latest/user-guide/install')
            sys.exit(1)
        self.conda = pathlib.Path(conda)
        self.this_dir = pathlib.Path(__file__).absolute().parent
        self.cenv_dir = self.this_dir / '.cenv'
        print('conda found at: {}'.format(self.conda))

        # reset custom config settings
        open(BAZELRC, 'w').close()

    def configure_conda(self, skip_cenv=None):
        if is_windows():
            env_file = self.this_dir / 'environment-windows.yml'
        else:
            env_file = self.this_dir / 'environment.yml'

        if self.cenv_dir.exists():
            print('Updating conda environment from: {}'.format(env_file))
            cmd = [
                str(self.conda),
                'env',
                'update',
                '-f',
                str(env_file),
                '-p',
                str(self.cenv_dir),
            ]
            subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL)
        else:
            if not skip_cenv:
                print('Creating conda environment from: {}'.format(env_file))
                cmd = [
                    str(self.conda),
                    'env',
                    'create',
                    '-f',
                    str(env_file),
                    '-p',
                    str(self.cenv_dir),
                ]
                subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL)

    def configure_cache(self, bazel_cache=None):
        print()
        if bazel_cache:
            print("Configuring bazel cache:", bazel_cache)
            write_to_bazelrc('build --remote_cache={}'.format(bazel_cache))
            if is_windows():
                pass  # TODO
            else:
                write_to_bazelrc('build --action_env={}'.format(UNIX_PATH))

    def configure_precommit(self, skip=False):
        if skip:
            return
        if is_windows():
            search_path = self.cenv_dir / 'Scripts'
        else:
            search_path = self.cenv_dir / 'bin'
        print()
        print('Searching for pre-commit in: {}'.format(search_path))
        pre_commit = shutil.which('pre-commit', path=str(search_path))
        if not pre_commit:
            print('pre-commit could not be found.')
            print('Is your conda environment created and up to date?')
            sys.exit(1)
        subprocess.run([pre_commit, 'install'], check=True)

    def configure_bazelisk(self, skip=False):
        if skip:
            return
        bazelisk = shutil.which('bazelisk')
        if not bazelisk:
            print('Please install bazelisk from:')
            print('https://github.com/bazelbuild/bazelisk')
            sys.exit(1)
        print()
        print('bazelisk version')
        subprocess.run([bazelisk, 'version'], check=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bazel_cache',
                        help='Enter a bazel cache server URL (omit argument to disable)')
    parser.add_argument('--ci_mode',
                        help='For CI machines, create conda env and skip bazelisk/precommit',
                        action='store_true')
    parser.add_argument('--skip_bazelisk',
                        help='Skip the bazelisk configuration step?',
                        action='store_true')
    parser.add_argument('--skip_conda_env',
                        help='Skip the conda environment creation step?',
                        action='store_true')
    parser.add_argument('--skip_precommit',
                        help='Skip the precommit configuration step?',
                        action='store_true')
    args = parser.parse_args()

    print("Configuring PlaidML build environment")
    print()

    if args.ci_mode:
        print("CI mode has been enabled. Overriding the following settings:")
        print("    --skip_bazelisk  = True")
        print("    --skip_conda_env = False")
        print("    --skip_precommit = True")
        print()

    cfg = Configure()
    cfg.configure_conda(False if args.ci_mode else args.skip_conda_env)
    cfg.configure_precommit(True if args.ci_mode else args.skip_precommit)
    cfg.configure_bazelisk(True if args.ci_mode else args.skip_bazelisk)
    cfg.configure_cache(args.bazel_cache)

    print()
    print("Your build is configured.")
    print("Use the following to run all unit tests:")
    print()
    print("bazelisk test //...")
    print()


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass