dev

#!/usr/bin/env python3
# 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], gdb: bool = False) -> None:
    """Run a command, terminating on failure."""
    cmd = [str(arg) for arg in command if arg is not None]
    if gdb:
        cmd = ["gdb", "-ex", "catch throw", "-ex", "run", "--args"] + cmd
    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 = ["tensor_tracker", "tests", "doc", "dev", "setup.py"]


@cli("-s", "--no-capture", action="store_false", dest="capture")
@cli("-k", "--filter")
def tests(capture: bool, filter: Optional[str]) -> None:
    """run Python tests"""
    run(
        [
            "python3",
            "-m",
            "pytest",
            "tests",
            None if filter else "--cov=tensor_tracker",
            *(["-k", filter] if filter else []),
            None if capture else "-s",
        ],
    )


@cli()
def lint() -> None:
    """run static analysis"""
    run(["python3", "-m", "flake8", *PYTHON_ROOTS])
    run(["python3", "-m", "mypy", *(r for r in PYTHON_ROOTS if r != "doc")])


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


@cli()
def copyright() -> None:
    """check for Graphcore copyright headers on relevant files"""
    command = (
        "find " + " ".join(PYTHON_ROOTS) + " -type f -not -name *.pyc -not -name *.png"
        " | 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", "doc/tensor_tracker"])
    run(
        [
            "python3",
            "-m",
            "pdoc",
            "--html",
            "--output-dir",
            "doc",
            "tensor_tracker",
        ]
    )
    for notebook in ["Example", "Usage"]:
        run(
            [
                "jupyter",
                "nbconvert",
                "--to",
                "html",
                f"doc/{notebook}.ipynb",
                "--output-dir",
                "doc/tensor_tracker",
                "--output",
                f"{notebook.lower()}.html",
            ]
        )


@cli("--skip", nargs="*", default=[], help="commands to skip")
def ci(skip: List[str] = []) -> None:
    """run all continuous integration tests & checks"""
    if "tests" not in skip:
        tests(capture=True, 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:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.set_defaults(action=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(action=value)

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


if __name__ == "__main__":
    _main()