ingest/fetch_sources.py

"""Fetch each book's repository and render it to markdown with mdBook.

mdBook defines these books' syntax, so we let it do the rendering. That resolves
the `{{#include}}` and `{{#rustdoc_include}}` directives exactly as the published
site does, and its markdown backend keeps the original fence attributes and the
`#` prefix on hidden lines.

We build each book with a minimal `book.toml` of our own rather than the
repository's, because theirs pull in preprocessors and link-checking backends that
need extra binaries and do not affect the text we want.

Requires the `mdbook` binary on PATH: https://rust-lang.github.io/mdBook/

    uv run python -m ingest.fetch_sources
    uv run python -m ingest.fetch_sources --book nomicon
    uv run python -m ingest.fetch_sources --refetch
"""

import argparse
import io
import json
import re
import shutil
import subprocess
import sys
import tarfile
from datetime import datetime, timezone
from pathlib import Path

import requests

from rag.config import BOOKS, MANIFEST_PATH, SOURCES_DIR, BookSpec
from rag.types import ManifestBook, Page

TARBALL_URL = "https://codeload.github.com/{repo}/tar.gz/{sha}"
REQUEST_TIMEOUT = 120

# Only what mdBook needs to render the text. `create-missing` is off to prevent
# generating unwritten chapters that might be referenced in SUMMARY.md.
BOOK_TOML = """[book]
src = "src"

[build]
create-missing = false

[output.markdown]
"""

# mdBook table of contents entries are indented markdown links, and the nesting is
# the topic structure. Numbered chapters are bullets, but prefix and suffix
# chapters carry none, so the marker is optional. Part headings are `# Title`
# lines between entries.
TOC_LINK = re.compile(r"^\s*(?:[-*]\s*)?\[(?P<title>[^\]]*)\]\((?P<href>[^)]*)\)")
PART_HEADING = re.compile(r"^#\s+(?P<title>.+?)\s*$")


def download_repo(session: requests.Session, spec: BookSpec, target: Path) -> None:
    response = session.get(
        TARBALL_URL.format(repo=spec.repo, sha=spec.sha), timeout=REQUEST_TIMEOUT
    )
    response.raise_for_status()
    with tarfile.open(fileobj=io.BytesIO(response.content), mode="r:gz") as archive:
        root = archive.getnames()[0].split("/")[0]
        archive.extractall(target.parent, filter="data")
    if target.exists():
        shutil.rmtree(target)
    (target.parent / root).rename(target)


def render_markdown(repo_dir: Path, out_dir: Path) -> int:
    if shutil.which("mdbook") is None:
        raise RuntimeError(
            "mdbook not found on PATH. Install it from https://rust-lang.github.io/mdBook/"
        )
    (repo_dir / "book.toml").write_text(BOOK_TOML)
    result = subprocess.run(
        ["mdbook", "build", "-d", str(out_dir)], cwd=repo_dir, capture_output=True, text=True
    )
    if result.returncode != 0:
        raise RuntimeError(f"mdbook build failed for {repo_dir.name}:\n{result.stderr.strip()}")

    # The markdown backend renders chapters only. SUMMARY.md is its input, and it is what gives us titles, ordering and part headings.
    shutil.copyfile(repo_dir / "src" / "SUMMARY.md", out_dir / "SUMMARY.md")
    return len(list(out_dir.rglob("*.md")))


def parse_summary(text: str, book: str, spec: BookSpec) -> list[Page]:
    pages: list[Page] = []
    part: str | None = None
    seen: set[str] = set()

    for line in text.splitlines():
        heading = PART_HEADING.match(line)
        if heading:
            part = heading.group("title")
            continue

        match = TOC_LINK.match(line)
        if not match:
            continue
        href = match.group("href").strip()
        if not href or part in spec.exclude_parts:
            continue

        # nomicon writes 19 of its entries as "./vec/vec-drain.md"
        path = href.split("#")[0].removeprefix("./")
        if not path.endswith(".md") or path in seen or path in spec.exclude_paths:
            continue
        seen.add(path)

        pages.append(
            {
                "book": book,
                "path": path,
                "title": match.group("title").strip(),
                "part": part,
                "url": f"{spec.base_url}{path[: -len('.md')]}.html",
            }
        )
    return pages


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--book", action="append", choices=sorted(BOOKS))
    parser.add_argument("--refetch", action="store_true", help="re-download and rebuild")
    args = parser.parse_args(argv[1:])

    fetched_at = datetime.now(timezone.utc).isoformat(timespec="seconds")
    books: dict[str, ManifestBook] = {}
    if MANIFEST_PATH.exists():
        books.update(json.loads(MANIFEST_PATH.read_text()).get("books", {}))

    session = requests.Session()
    session.headers["User-Agent"] = "rust-docs-rag ingestion (rust-lang docs, educational use)"

    for book in args.book or sorted(BOOKS):
        spec = BOOKS[book]
        out_dir, repo_dir = SOURCES_DIR / book, SOURCES_DIR / f"{book}.repo"
        print(f"{spec.title} ({spec.repo} @ {spec.sha})")

        if out_dir.exists() and not args.refetch:
            print("  cached")
        else:
            out_dir.parent.mkdir(parents=True, exist_ok=True)
            download_repo(session, spec, repo_dir)
            print(f"  rendered {render_markdown(repo_dir, out_dir)} markdown files")
            shutil.rmtree(repo_dir)

        pages = parse_summary((out_dir / "SUMMARY.md").read_text(encoding="utf-8"), book, spec)
        present = [p for p in pages if (out_dir / p["path"]).exists()]
        print(f"  {len(pages)} pages in SUMMARY.md, {len(pages) - len(present)} not rendered")

        books[book] = {
            "title": spec.title,
            "repo": spec.repo,
            "sha": spec.sha,
            "base_url": spec.base_url,
            "license": spec.license,
            "excluded_parts": list(spec.exclude_parts),
            "page_count": len(present),
            "pages": present,
        }

    MANIFEST_PATH.parent.mkdir(parents=True, exist_ok=True)
    manifest = {"fetched_at": fetched_at, "books": books}
    MANIFEST_PATH.write_text(json.dumps(manifest, indent=2) + "\n")
    total = sum(b["page_count"] for b in books.values())
    print(f"\n{total} pages across {len(books)} books -> {MANIFEST_PATH.name}")
    return 0


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