run_dev.py

#!/usr/bin/env python3
"""根目录一键启动开发环境。"""

from __future__ import annotations

import os
import shutil
import signal
import subprocess
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path


ROOT_DIR = Path(__file__).resolve().parent
BACKEND_DIR = ROOT_DIR / "backend"
FRONTEND_DIR = ROOT_DIR / "frontend"


@dataclass(slots=True)
class ManagedProcess:
    """受管进程。"""

    name: str
    process: subprocess.Popen[str]


def find_command(name: str) -> str:
    """查找命令路径。"""
    path = shutil.which(name)
    if path:
        return path
    raise SystemExit(f"缺少命令：{name}，请先安装后再运行。")


def run_blocking_step(name: str, args: list[str], cwd: Path, env: dict[str, str] | None = None) -> None:
    """执行阻塞式准备步骤。"""
    print(f"[setup] {name}：{' '.join(args)}")
    completed = subprocess.run(
        args,
        cwd=str(cwd),
        env=env,
        check=False,
    )
    if completed.returncode != 0:
        raise SystemExit(completed.returncode)


def ensure_backend_ready(uv_cmd: str) -> None:
    """确保后端虚拟环境可用。"""
    if not BACKEND_DIR.exists():
        raise SystemExit(f"后端目录不存在：{BACKEND_DIR}")
    if not (BACKEND_DIR / ".venv").exists():
        run_blocking_step("初始化后端依赖", [uv_cmd, "sync"], BACKEND_DIR, env=os.environ.copy())


def ensure_frontend_ready(pnpm_cmd: str) -> None:
    """确保前端依赖可用。"""
    if not FRONTEND_DIR.exists():
        raise SystemExit(f"前端目录不存在：{FRONTEND_DIR}")
    if not (FRONTEND_DIR / "node_modules").exists():
        env = os.environ.copy()
        env["COREPACK_ENABLE_STRICT"] = "0"
        run_blocking_step("初始化前端依赖", [pnpm_cmd, "install"], FRONTEND_DIR, env=env)


def stream_output(name: str, process: subprocess.Popen[str]) -> None:
    """实时打印子进程输出。"""
    assert process.stdout is not None
    for line in process.stdout:
        print(f"[{name}] {line}", end="")


def start_process(name: str, args: list[str], cwd: Path, env: dict[str, str] | None = None) -> ManagedProcess:
    """启动子进程并挂接输出。"""
    kwargs: dict[str, object] = {
        "args": args,
        "cwd": str(cwd),
        "env": env,
        "stderr": subprocess.STDOUT,
        "stdout": subprocess.PIPE,
        "text": True,
        "bufsize": 1,
    }
    if os.name == "nt":
        kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP  # type: ignore[attr-defined]
    else:
        kwargs["preexec_fn"] = os.setsid
    process = subprocess.Popen(**kwargs)  # type: ignore[arg-type]
    threading.Thread(target=stream_output, args=(name, process), daemon=True).start()
    return ManagedProcess(name=name, process=process)


def terminate_process(item: ManagedProcess) -> None:
    """终止单个子进程。"""
    process = item.process
    if process.poll() is not None:
        return
    try:
        if os.name == "nt":
            process.terminate()
        else:
            os.killpg(os.getpgid(process.pid), signal.SIGTERM)
    except ProcessLookupError:
        return


def kill_process(item: ManagedProcess) -> None:
    """强制杀掉单个子进程。"""
    process = item.process
    if process.poll() is not None:
        return
    try:
        if os.name == "nt":
            process.kill()
        else:
            os.killpg(os.getpgid(process.pid), signal.SIGKILL)
    except ProcessLookupError:
        return


def shutdown(processes: list[ManagedProcess]) -> None:
    """依次关闭所有子进程。"""
    for item in reversed(processes):
        terminate_process(item)
    deadline = time.time() + 5
    while time.time() < deadline:
        alive = [item for item in processes if item.process.poll() is None]
        if not alive:
            return
        time.sleep(0.2)
    for item in reversed(processes):
        kill_process(item)


def main() -> int:
    """启动后端 API、Worker 和前端开发服务。"""
    uv_cmd = find_command("uv")
    pnpm_cmd = find_command("pnpm")

    ensure_backend_ready(uv_cmd)
    ensure_frontend_ready(pnpm_cmd)

    print()
    print("开发环境启动中：")
    print(f"- 后端 API：{BACKEND_DIR}")
    print(f"- Worker：{BACKEND_DIR}")
    print(f"- 前端：{FRONTEND_DIR}")
    print()
    print("访问地址：")
    print("- 前端：http://127.0.0.1:5666")
    print("- 后端：http://127.0.0.1:8000")
    print()
    print("按 Ctrl+C 可一次性关闭全部进程。")
    print()

    backend_env = os.environ.copy()
    backend_env["PYTHONUNBUFFERED"] = "1"
    frontend_env = os.environ.copy()
    frontend_env["COREPACK_ENABLE_STRICT"] = "0"

    processes = [
        start_process(
            "backend-api",
            [uv_cmd, "run", "uvicorn", "app.main:app", "--reload", "--host", "127.0.0.1", "--port", "8000"],
            BACKEND_DIR,
            env=backend_env,
        ),
        start_process(
            "worker",
            [uv_cmd, "run", "python", "-m", "app.worker.main"],
            BACKEND_DIR,
            env=backend_env,
        ),
        start_process(
            "frontend",
            [pnpm_cmd, "-F", "@vben/web-antd", "dev"],
            FRONTEND_DIR,
            env=frontend_env,
        ),
    ]

    stop_requested = False

    def handle_signal(signum: int, _frame) -> None:
        nonlocal stop_requested
        if stop_requested:
            return
        stop_requested = True
        print()
        print(f"收到信号 {signum}，开始关闭全部进程...")
        shutdown(processes)

    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    try:
        while True:
            for item in processes:
                code = item.process.poll()
                if code is None:
                    continue
                print()
                print(f"{item.name} 已退出，退出码：{code}")
                shutdown(processes)
                return code
            time.sleep(0.5)
    except KeyboardInterrupt:
        handle_signal(signal.SIGINT, None)
        return 130


if __name__ == "__main__":
    sys.exit(main())