polyester-sdk 0.1.0-alpha.26

#!/usr/bin/env python3
"""Regenerate src/proto and src/connect_gen module trees from src/gen/*."""
from __future__ import annotations

from pathlib import Path
import shutil

ROOT = Path(__file__).resolve().parents[1]
KEYWORDS = {
    "type", "mod", "pub", "fn", "struct", "enum", "use", "crate", "self",
    "super", "async", "await", "match", "where", "impl", "move", "ref",
    "mut", "const", "static", "trait", "unsafe", "extern", "dyn",
}


def rust_ident(part: str) -> str:
    return f"r#{part}" if part in KEYWORDS else part


def parse_pkg(filename: str) -> list[str]:
    assert filename.endswith(".rs")
    return filename[:-3].split(".")


class Node:
    __slots__ = ("file", "children")

    def __init__(self) -> None:
        self.file: Path | None = None
        self.children: dict[str, Node] = {}


def build_tree(files: list[Path]) -> Node:
    root = Node()
    for f in files:
        node = root
        for p in parse_pkg(f.name):
            node = node.children.setdefault(p, Node())
        if node.file is not None:
            raise SystemExit(f"duplicate package for {f}")
        node.file = f
    return root


def write_tree(out_dir: Path, node: Node, depth: int, gen_kind: str) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    lines = [
        "// @generated by scripts/gen_module_tree.py — DO NOT EDIT BY HAND.",
        "#![allow(unused_imports, dead_code, clippy::all, warnings, unused_qualifications)]",
        "",
    ]
    if node.file is not None:
        include_path = "/".join([".."] * (depth + 1) + ["gen", gen_kind, node.file.name])
        lines.append(f'include!("{include_path}");')
        lines.append("")
    for name in sorted(node.children):
        child = node.children[name]
        ident = rust_ident(name)
        write_tree(out_dir / name, child, depth + 1, gen_kind)
        if ident.startswith("r#"):
            lines.append(f'#[path = "{name}/mod.rs"]')
        lines.append(f"pub mod {ident};")
    (out_dir / "mod.rs").write_text("\n".join(lines) + "\n")


SKIP_PACKAGES = {
    "gnostic.openapi.v3.rs",
    "buf.validate.rs",
    "google.api.rs",
}


def main() -> None:
    for out, kind in (
        (ROOT / "src/proto", "buffa"),
        (ROOT / "src/connect_gen", "connect"),
    ):
        if out.exists():
            shutil.rmtree(out)
        gen = ROOT / "src/gen" / kind
        files = [p for p in sorted(gen.glob("*.rs")) if p.name not in SKIP_PACKAGES]
        write_tree(out, build_tree(files), 0, kind)
        print(f"wrote {out}")


if __name__ == "__main__":
    main()