stack/dev/scripts/pkg_inventory.py

"""Scan the repo for all apt, apk, and PyPI packages in use.

Parses Dockerfiles, pyproject.toml, CI workflows, and shell scripts to
produce a canonical JSON manifest at data/pkg-manifest.json.

Usage:
    uv run python dev/scripts/pkg_inventory.py
    uv run python dev/scripts/pkg_inventory.py --check  # exit 1 if manifest is stale
"""

from __future__ import annotations

import json
import re
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
MANIFEST_PATH = ROOT / "data" / "pkg-manifest.json"

# Directories to skip when scanning Dockerfiles
SKIP_DIRS = {"node_modules", ".git", "__pycache__", "plugins"}


def _drop_gitignored(paths: list[Path]) -> list[Path]:
    """Filter out gitignored paths (e.g. vendored checkouts like
    infra/marimo/src) so local manifests match CI, where ignored
    files are absent from the checkout."""
    if not paths:
        return paths
    try:
        result = subprocess.run(
            ["git", "-C", str(ROOT), "check-ignore", "--stdin"],
            input="\n".join(str(p) for p in paths),
            capture_output=True,
            text=True,
            timeout=30,
        )
    except (OSError, subprocess.TimeoutExpired):
        return paths
    # exit 0: some ignored, 1: none ignored, 128: fatal (not a repo)
    if result.returncode not in (0, 1):
        return paths
    ignored = set(result.stdout.splitlines())
    return [p for p in paths if str(p) not in ignored]


# ---------------------------------------------------------------------------
# Dockerfile parsers
# ---------------------------------------------------------------------------


def _join_continuation_lines(text: str) -> str:
    """Merge backslash-continued lines into single lines."""
    return re.sub(r"\\\s*\n\s*", " ", text)


def _parse_apt_packages(text: str) -> list[str]:
    """Extract packages from apt-get install commands."""
    text = _join_continuation_lines(text)
    pkgs: set[str] = set()
    for m in re.finditer(r"apt-get\s+install\s+(?:-\S+\s+)*(.+?)(?:&&|;|\n|$)", text):
        tokens = m.group(1).split()
        for tok in tokens:
            tok = tok.strip()
            if tok and not tok.startswith("-") and not tok.startswith("#"):
                # strip version pin like =1.2.3
                name = re.split(r"[=<>]", tok)[0]
                if name and not name.startswith("/"):
                    pkgs.add(name)
    return sorted(pkgs)


def _parse_apk_packages(text: str) -> list[str]:
    """Extract packages from apk add commands."""
    text = _join_continuation_lines(text)
    pkgs: set[str] = set()
    for m in re.finditer(r"apk\s+add\s+(?:-\S+\s+)*(.+?)(?:&&|;|\n|$)", text):
        tokens = m.group(1).split()
        for tok in tokens:
            tok = tok.strip()
            if tok and not tok.startswith("-") and not tok.startswith("#"):
                name = re.split(r"[=<>~]", tok)[0]
                if name:
                    pkgs.add(name)
    return sorted(pkgs)


def _parse_uv_add_packages(text: str) -> list[dict]:
    """Extract packages from 'uv add' commands in Dockerfiles."""
    text = _join_continuation_lines(text)
    pkgs: list[dict] = []
    for m in re.finditer(r"uv\s+add\s+(.+?)(?:&&|;|\n|$)", text):
        tokens = m.group(1).split()
        for tok in tokens:
            tok = tok.strip().strip('"').strip("'")
            if tok.startswith("-"):
                continue
            if tok.startswith("http"):
                continue
            if tok:
                pkgs.append(_parse_pypi_spec(tok))
    return pkgs


def _parse_pip_install_packages(text: str) -> list[dict]:
    """Extract packages from 'pip install' commands in Dockerfiles."""
    text = _join_continuation_lines(text)
    pkgs: list[dict] = []
    for m in re.finditer(r"pip\s+install\s+(.+?)(?:&&|;|\n|$)", text):
        tokens = m.group(1).split()
        for tok in tokens:
            tok = tok.strip().strip('"').strip("'")
            if tok.startswith("-"):
                continue
            if tok and tok != ".":
                pkgs.append(_parse_pypi_spec(tok))
    return pkgs


def _parse_uv_run_with_packages(text: str) -> list[dict]:
    """Extract packages from 'uv run --with pkg' commands."""
    text = _join_continuation_lines(text)
    pkgs: list[dict] = []
    for m in re.finditer(r"uv\s+run\s+--with\s+(\S+)", text):
        tok = m.group(1).strip('"').strip("'")
        if tok:
            pkgs.append(_parse_pypi_spec(tok))
    return pkgs


def _parse_pypi_spec(spec: str) -> dict:
    """Parse a PEP 508 spec like 'narwhals>=2.17.0' or 'marimo[recommended]'."""
    # strip extras
    extras = ""
    if "[" in spec:
        base, rest = spec.split("[", 1)
        extras = rest.split("]")[0]
        spec = base + rest.split("]")[-1] if "]" in rest else base
    m = re.match(r"([a-zA-Z0-9_-]+)(.*)", spec)
    if not m:
        return {"name": spec, "version": "", "extras": ""}
    return {
        "name": m.group(1).lower().replace("_", "-"),
        "version": m.group(2).strip(),
        "extras": extras,
    }


def _parse_base_images(text: str) -> list[str]:
    """Extract FROM base images from Dockerfiles."""
    images: list[str] = []
    for m in re.finditer(r"^FROM\s+(\S+)", text, re.MULTILINE):
        img = m.group(1)
        if img not in images:
            images.append(img)
    return images


# ---------------------------------------------------------------------------
# CI workflow parser
# ---------------------------------------------------------------------------


def _parse_ci_curl_tools(text: str) -> list[dict]:
    """Extract tools installed via curl in CI workflows."""
    tools: list[dict] = []
    patterns = [
        (r"astral\.sh/uv/install\.sh", "uv", "latest"),
        (r"go-containerregistry.*crane", "crane", "latest"),
        (r"aquasecurity/trivy", "trivy", "latest"),
    ]
    for pattern, name, version in patterns:
        if re.search(pattern, text):
            tools.append({"name": name, "version": version})
    return tools


# ---------------------------------------------------------------------------
# pyproject.toml parser (simple, no toml dependency)
# ---------------------------------------------------------------------------


def _is_self_ref(spec: str) -> bool:
    """Check if a dependency spec is a self-reference (e.g. stack[aco])."""
    name = re.split(r"[\[>=<~!]", spec)[0].strip().lower().replace("_", "-")
    return name == "stack"


def _extract_dep_specs(block: str) -> list[str]:
    """Extract dependency spec strings from a TOML array body.

    Pulls the first quoted string from each non-comment line. This
    correctly skips inline comments (e.g. ``"pymupdf>=1.24",  # AGPL``),
    trailing commas, and surrounding whitespace — the previous chained
    ``.strip()`` approach broke on inline comments.
    """
    specs: list[str] = []
    for line in block.splitlines():
        s = line.strip()
        if not s or s.startswith("#"):
            continue
        m = re.match(r'["\']([^"\']+)["\']', s)
        if m:
            specs.append(m.group(1).strip())
    return specs


def _parse_pyproject_deps(text: str) -> tuple[list[dict], list[dict]]:
    """Parse dependencies from pyproject.toml without a TOML library.

    Extracts:
    - [project] dependencies (prod)
    - [project.optional-dependencies] all extras (prod)
    - [dependency-groups] dev (dev)
    - [build-system] requires (prod)

    Filters out self-references (stack[*]) which can't be downloaded.
    """
    prod: list[dict] = []
    dev: list[dict] = []

    # prod deps — [project] dependencies
    m = re.search(r"^dependencies\s*=\s*\[(.*?)\]", text, re.MULTILINE | re.DOTALL)
    if m:
        for spec in _extract_dep_specs(m.group(1)):
            if not _is_self_ref(spec):
                prod.append(_parse_pypi_spec(spec))

    # optional deps — [project.optional-dependencies] all sections
    # Extract the section content, then parse each "name = [...]" block.
    # Can't use simple .*? because extras brackets (e.g. stack[conf])
    # contain ] which terminates the match early.
    opt_section = re.search(
        r"\[project\.optional-dependencies\](.*?)(?:\n\[|\Z)",
        text,
        re.DOTALL,
    )
    if opt_section:
        # Match each key = [ ... ] block by finding balanced brackets
        for m in re.finditer(
            r"^(\w+)\s*=\s*\[",
            opt_section.group(1),
            re.MULTILINE,
        ):
            start = m.end()
            rest = opt_section.group(1)[start:]
            # Find closing ] that isn't inside a string bracket
            depth = 1
            i = 0
            while i < len(rest) and depth > 0:
                if rest[i] == "[":
                    depth += 1
                elif rest[i] == "]":
                    depth -= 1
                i += 1
            block = rest[: i - 1] if depth == 0 else rest
            for spec in _extract_dep_specs(block):
                if not _is_self_ref(spec):
                    prod.append(_parse_pypi_spec(spec))

    # dev deps — [dependency-groups] dev
    m = re.search(r"dev\s*=\s*\[(.*?)\]", text, re.MULTILINE | re.DOTALL)
    if m:
        for spec in _extract_dep_specs(m.group(1)):
            if not _is_self_ref(spec):
                dev.append(_parse_pypi_spec(spec))

    # build-system requires
    m = re.search(
        r"\[build-system\].*?requires\s*=\s*\[(.*?)\]",
        text,
        re.MULTILINE | re.DOTALL,
    )
    if m:
        for spec in _extract_dep_specs(m.group(1)):
            if not _is_self_ref(spec):
                prod.append(_parse_pypi_spec(spec))

    return prod, dev


# ---------------------------------------------------------------------------
# pnpm / Node parser
# ---------------------------------------------------------------------------


def _parse_package_json_deps(text: str) -> list[dict]:
    """Parse npm dependencies from package.json."""
    data = json.loads(text)
    pkgs: list[dict] = []
    for section in ("dependencies", "devDependencies"):
        for name, version in data.get(section, {}).items():
            pkgs.append({"name": name, "version": version})
    return pkgs


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def scan() -> dict:
    """Scan repo and return the package manifest."""
    manifest: dict = {
        "apt": {},
        "apk": {},
        "pypi": {
            "project_prod": [],
            "project_dev": [],
            "notebook": [],
            "dockerfile_adhoc": [],
        },
        "npm": [],
        "ci_tools": [],
        "base_images": {},
    }

    # --- Dockerfiles ---
    dockerfiles = list(ROOT.glob("**/Dockerfile")) + list(ROOT.glob("**/Dockerfile.*"))
    dockerfiles = [
        df for df in dockerfiles if not any(skip in df.parts for skip in SKIP_DIRS)
    ]
    dockerfiles = _drop_gitignored(dockerfiles)
    for df in sorted(dockerfiles):
        rel = str(df.relative_to(ROOT))
        text = df.read_text()

        apt = _parse_apt_packages(text)
        if apt:
            manifest["apt"][rel] = apt

        apk = _parse_apk_packages(text)
        if apk:
            manifest["apk"][rel] = apk

        images = _parse_base_images(text)
        if images:
            manifest["base_images"][rel] = images

        # uv add in Dockerfiles (notebook pattern)
        uv_pkgs = _parse_uv_add_packages(text)
        if uv_pkgs:
            manifest["pypi"]["notebook"].extend(uv_pkgs)

        # pip install in Dockerfiles
        pip_pkgs = _parse_pip_install_packages(text)
        if pip_pkgs:
            manifest["pypi"]["dockerfile_adhoc"].extend(pip_pkgs)

        # uv run --with
        with_pkgs = _parse_uv_run_with_packages(text)
        if with_pkgs:
            manifest["pypi"]["dockerfile_adhoc"].extend(with_pkgs)

    # --- pyproject.toml ---
    pyproject = ROOT / "pyproject.toml"
    if pyproject.exists():
        prod, dev = _parse_pyproject_deps(pyproject.read_text())
        manifest["pypi"]["project_prod"] = prod
        manifest["pypi"]["project_dev"] = dev

    # --- package.json (docs) ---
    pkg_json = ROOT / "docs" / "package.json"
    if pkg_json.exists():
        manifest["npm"] = _parse_package_json_deps(pkg_json.read_text())

    # --- CI workflows ---
    ci_tools_seen: set[str] = set()
    for wf in sorted((ROOT / ".gitea" / "workflows").glob("*.yml")):
        text = wf.read_text()
        for tool in _parse_ci_curl_tools(text):
            if tool["name"] not in ci_tools_seen:
                ci_tools_seen.add(tool["name"])
                manifest["ci_tools"].append(tool)

    # --- Deduplicate & sort ---
    for section in ("notebook", "dockerfile_adhoc"):
        seen: set[str] = set()
        deduped: list[dict] = []
        for pkg in manifest["pypi"][section]:
            if pkg["name"] not in seen:
                seen.add(pkg["name"])
                deduped.append(pkg)
        manifest["pypi"][section] = sorted(deduped, key=lambda p: p["name"])

    return manifest


def main() -> None:
    manifest = scan()
    output = json.dumps(manifest, indent=2, sort_keys=False) + "\n"

    if "--check" in sys.argv:
        if MANIFEST_PATH.exists():
            existing = MANIFEST_PATH.read_text()
            if existing == output:
                print("pkg-manifest.json is up to date.")
                raise SystemExit(0)
            else:
                print("pkg-manifest.json is STALE — re-run without --check.")
                raise SystemExit(1)
        else:
            print("pkg-manifest.json does not exist — run without --check first.")
            raise SystemExit(1)

    MANIFEST_PATH.write_text(output)
    # Summary
    apt_count = sum(len(v) for v in manifest["apt"].values())
    apk_count = sum(len(v) for v in manifest["apk"].values())
    pypi_count = sum(
        len(manifest["pypi"][k])
        for k in ("project_prod", "project_dev", "notebook", "dockerfile_adhoc")
    )
    npm_count = len(manifest["npm"])
    ci_count = len(manifest["ci_tools"])
    img_count = sum(len(v) for v in manifest["base_images"].values())
    print(f"Wrote {MANIFEST_PATH.relative_to(ROOT)}")
    print(
        f"  apt: {apt_count}  apk: {apk_count}  pypi: {pypi_count}"
        f"  npm: {npm_count}  ci_tools: {ci_count}  base_images: {img_count}"
    )


if __name__ == "__main__":
    main()