server.py

"""Serve the codebase graph as JSON to a static frontend.

Usage: python server.py <path_to_python_codebase> [--exclude DIR ...] [--port PORT]
"""

import argparse
import sys
from pathlib import Path

import uvicorn
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

from walker import walk_python_files
from parser import parse_codebase
from graph import Graph, build_graph


STATIC_DIR = Path(__file__).parent / "static"


def build_codebase_graph(codebase_path: Path, exclude: set[str]) -> tuple[Graph, dict]:
    root = codebase_path.resolve()
    files = list(walk_python_files(root, extra_skip=exclude))
    if not files:
        raise SystemExit(f"no python files found under {root}")
    parsed = parse_codebase(root, files)
    failed = [p for p in parsed if p.parse_error]
    for p in failed:
        print(f"warn: {p.relpath}: {p.parse_error}", file=sys.stderr)
    g = build_graph(parsed)
    meta = {
        "root": str(root),
        "file_count": len(parsed),
        "failed_count": len(failed),
        "node_count": len(g.nodes),
        "edge_count": len(g.edges),
    }
    print(f"loaded {meta['file_count']} files ({meta['failed_count']} failed), "
          f"{meta['node_count']} nodes, {meta['edge_count']} edges", file=sys.stderr)
    return g, meta


def create_app(codebase_path: Path, exclude: set[str]) -> FastAPI:
    g, meta = build_codebase_graph(codebase_path, exclude)
    payload = {**g.to_dict(), "meta": meta}

    app = FastAPI(title="codemap")

    @app.middleware("http")
    async def no_cache(request, call_next):
        response = await call_next(request)
        response.headers["Cache-Control"] = "no-store, must-revalidate"
        response.headers["Pragma"] = "no-cache"
        return response

    @app.get("/graph.json", response_class=JSONResponse)
    def graph_json() -> dict:
        return payload

    app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="static")
    return app


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("path", help="path to a Python codebase")
    ap.add_argument("--exclude", action="append", default=[],
                    help="directory name to skip (can be passed multiple times)")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8000)
    args = ap.parse_args()
    app = create_app(Path(args.path), set(args.exclude))
    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")


if __name__ == "__main__":
    main()