tasks.py

"""Simple task runner for common project workflows."""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
import time


def run_cmd(cmd: list[str]) -> int:
    print(">", " ".join(cmd))
    return subprocess.call(cmd)


def task_install() -> int:
    return run_cmd([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])


def task_run() -> int:
    return run_cmd([sys.executable, "-m", "uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "8000"])


def task_space() -> int:
    return run_cmd([sys.executable, "app.py"])


def task_smoke() -> int:
    return run_cmd([sys.executable, "smoke_test.py"])


def task_compile() -> int:
    return run_cmd([sys.executable, "-m", "compileall", "."])


def task_docker_up() -> int:
    return run_cmd(["docker", "compose", "up", "--build"])


def task_docker_down() -> int:
    return run_cmd(["docker", "compose", "down"])


def task_hf_upload(repo_id: str, token: str, private: bool) -> int:
    cmd = [sys.executable, "upload_to_hf.py", "--repo-id", repo_id, "--token", token]
    if private:
        cmd.append("--private")
    return run_cmd(cmd)


def task_serve_smoke() -> int:
    """Start API, run smoke test, then stop API."""
    env = os.environ.copy()
    server_cmd = [sys.executable, "-m", "uvicorn", "api.main:app", "--host", "127.0.0.1", "--port", "8000"]
    print(">", " ".join(server_cmd))
    server = subprocess.Popen(server_cmd, env=env)
    try:
        time.sleep(3)
        return run_cmd([sys.executable, "smoke_test.py"])
    finally:
        server.terminate()
        try:
            server.wait(timeout=10)
        except subprocess.TimeoutExpired:
            server.kill()


def main():
    parser = argparse.ArgumentParser(description="Project task runner")
    sub = parser.add_subparsers(dest="task", required=True)

    sub.add_parser("install")
    sub.add_parser("run")
    sub.add_parser("space")
    sub.add_parser("smoke")
    sub.add_parser("compile")
    sub.add_parser("docker-up")
    sub.add_parser("docker-down")
    sub.add_parser("serve-smoke")
    hf = sub.add_parser("hf-upload")
    hf.add_argument("--repo-id", required=True)
    hf.add_argument("--token", required=True)
    hf.add_argument("--private", action="store_true")

    args = parser.parse_args()

    if args.task == "install":
        code = task_install()
    elif args.task == "run":
        code = task_run()
    elif args.task == "space":
        code = task_space()
    elif args.task == "smoke":
        code = task_smoke()
    elif args.task == "compile":
        code = task_compile()
    elif args.task == "docker-up":
        code = task_docker_up()
    elif args.task == "docker-down":
        code = task_docker_down()
    elif args.task == "serve-smoke":
        code = task_serve_smoke()
    else:
        code = task_hf_upload(args.repo_id, args.token, args.private)

    raise SystemExit(code)


if __name__ == "__main__":
    main()