dev

#!/usr/bin/env python
# Copyright (c) 2023 Graphcore Ltd. All rights reserved.

"""Dev task launcher."""

import argparse
import datetime
import os
import subprocess
import sys
from pathlib import Path
from typing import Any, Callable, Iterable, List, Optional, TypeVar

# Utilities


def run(command: Iterable[Any]) -> None:
    """Run a command, terminating on failure."""
    cmd = [str(arg) for arg in command if arg is not None]
    print("$ " + " ".join(cmd), file=sys.stderr)
    environ = os.environ.copy()
    environ["PYTHONPATH"] = f"{os.getcwd()}:{environ.get('PYTHONPATH', '')}"
    exit_code = subprocess.call(cmd, env=environ)
    if exit_code:
        sys.exit(exit_code)


T = TypeVar("T")


def cli(*args: Any, **kwargs: Any) -> Callable[[T], T]:
    """Declare a CLI command / arguments for that command."""

    def wrap(func: T) -> T:
        if not hasattr(func, "cli_args"):
            setattr(func, "cli_args", [])
        if args or kwargs:
            getattr(func, "cli_args").append((args, kwargs))
        return func

    return wrap


# Commands

PYTHON_ROOTS = ["unit_scaling", "dev", "examples"]


@cli("-k", "--filter")
def tests(filter: Optional[str]) -> None:
    """run Python tests"""
    run(
        [
            "python",
            "-m",
            "pytest",
            "unit_scaling",
            None if filter else "--cov=unit_scaling",
            *(["-k", filter] if filter else []),
        ]
    )


@cli("commands", nargs="*")
def python(commands: List[Any]) -> None:
    """run Python with the current directory on PYTHONPATH, for development"""
    run(["python"] + commands)


@cli()
def lint() -> None:
    """run static analysis"""
    run(["python", "-m", "flake8", *PYTHON_ROOTS])
    run(["python", "-m", "mypy", *PYTHON_ROOTS])


@cli("--check", action="store_true")
def format(check: bool) -> None:
    """autoformat all sources"""
    run(["python", "-m", "black", "--check" if check else None, *PYTHON_ROOTS])
    run(["python", "-m", "isort", "--check" if check else None, *PYTHON_ROOTS])


@cli()
def copyright() -> None:
    """check for Graphcore copyright headers on relevant files"""
    command = (
        f"find {' '.join(PYTHON_ROOTS)} -type f -not -name *.pyc"
        " | xargs grep -L 'Copyright (c) 202. Graphcore Ltd[.] All rights reserved[.]'"
    )
    print(f"$ {command}", file=sys.stderr)
    # Note: grep exit codes are not consistent between versions, so we don't use
    # check=True
    output = (
        subprocess.run(
            command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
        )
        .stdout.decode()
        .strip()
    )
    if output:
        print(
            "Error - failed copyright header check in:\n   "
            + output.replace("\n", "\n   "),
            file=sys.stderr,
        )
        print("Template(s):")
        comment_prefixes = {
            {".cpp": "//"}.get(Path(f).suffix, "#") for f in output.split("\n")
        }
        for prefix in comment_prefixes:
            print(
                (
                    f"{prefix} Copyright (c) {datetime.datetime.now().year}"
                    " Graphcore Ltd. All rights reserved."
                ),
                file=sys.stderr,
            )
        sys.exit(1)


@cli()
def doc() -> None:
    """generate API documentation"""
    subprocess.call(["rm", "-r", "docs/generated", "docs/_build"])
    run(
        [
            "make",
            "-C",
            "docs",
            "html",
        ]
    )


@cli(
    "-s",
    "--skip",
    nargs="*",
    default=[],
    choices=["tests", "lint", "format", "copyright"],
    help="commands to skip",
)
def ci(skip: List[str] = []) -> None:
    """run all continuous integration tests & checks"""
    if "tests" not in skip:
        tests(filter=None)
    if "lint" not in skip:
        lint()
    if "format" not in skip:
        format(check=True)
    if "copyright" not in skip:
        copyright()
    if "doc" not in skip:
        doc()


# Script


def _main() -> None:
    # Build an argparse command line by finding globals in the current module
    # that are marked via the @cli() decorator. Each one becomes a subcommand
    # running that function, usage "$ ./dev fn_name ...args".
    parser = argparse.ArgumentParser(description=__doc__)
    parser.set_defaults(command=ci)

    subs = parser.add_subparsers()
    for key, value in globals().items():
        if hasattr(value, "cli_args"):
            sub = subs.add_parser(key.replace("_", "-"), help=value.__doc__)
            for args, kwargs in value.cli_args:
                sub.add_argument(*args, **kwargs)
            sub.set_defaults(command=value)

    cli_args = vars(parser.parse_args())
    command = cli_args.pop("command")
    command(**cli_args)


if __name__ == "__main__":
    _main()