md_indexer.py

#!/usr/bin/env python3
"""Create JSON indexes from Markdown front matter and ATX headings.

Usage: python3.14 md_indexer.py PATH [...] [-r] [-o OUTPUT_DIR]
PyYAML is optional; without it, front matter is stored as raw text.
"""

from __future__ import annotations

import argparse
import json
import re
from collections.abc import Iterable
from pathlib import Path
from typing import Any, BinaryIO

try:
    import yaml
except ImportError:  # pragma: no cover - depends on local environment
    yaml = None


HEADING_RE = re.compile(rb"^(#{1,6})(?:[ \t]+|$)(.*?)[ \t]*#*[ \t]*$")
FENCE_RE = re.compile(rb"^[ \t]{0,3}(`{3,}|~{3,})")
WORD_RE = re.compile(r"[\w]+(?:[-'][\w]+)*", re.UNICODE)


def decode_text(value: bytes | bytearray, encoding: str) -> str:
    return value.decode(encoding, errors="replace")


def strip_heading_text(raw_text: bytes, encoding: str) -> str:
    text = decode_text(raw_text.strip(), encoding).strip()
    if text.endswith("#"):
        text = re.sub(r"[ \t]+#+$", "", text).strip()
    return text


def count_words(text: str) -> int:
    return len(WORD_RE.findall(text))


def count_words_in_bytes(value: bytes, encoding: str) -> int:
    return count_words(decode_text(value, encoding))


def parse_frontmatter(handle: BinaryIO, encoding: str) -> tuple[dict[str, Any] | None, int, int]:
    handle.seek(0)
    first_line = handle.readline()
    if first_line.rstrip(b"\r\n") != b"---":
        handle.seek(0)
        return None, 0, 1

    raw_yaml = bytearray()
    line_number = 2

    while True:
        line = handle.readline()
        if not line:
            handle.seek(0)
            return None, 0, 1

        stripped_line = line.rstrip(b"\r\n").strip()
        next_offset = handle.tell()
        if stripped_line in {b"---", b"..."}:
            if yaml is None:
                frontmatter: dict[str, Any] = {
                    "_error": "PyYAML is not installed; raw frontmatter was not parsed.",
                    "_raw": decode_text(raw_yaml, encoding),
                }
            else:
                try:
                    parsed = yaml.safe_load(decode_text(raw_yaml, encoding))
                    frontmatter = parsed if isinstance(parsed, dict) else {}
                except yaml.YAMLError as error:
                    frontmatter = {
                        "_error": str(error),
                        "_raw": decode_text(raw_yaml, encoding),
                    }
            return frontmatter, next_offset, line_number + 1

        raw_yaml.extend(line)
        line_number += 1


def json_default(value: Any) -> str:
    return str(value)


def make_heading(
    level: int,
    title: str,
    line_number: int,
    heading_byte_start: int,
    content_byte_start: int,
) -> dict[str, Any]:
    return {
        "level": level,
        "title": title,
        "line": line_number,
        "heading_byte_start": heading_byte_start,
        "content_byte_start": content_byte_start,
        "content_byte_length": 0,
        "content_line_count": 0,
        "content_word_count": 0,
        "children": [],
    }


def close_heading(heading: dict[str, Any], end_byte: int, end_line: int) -> None:
    heading["content_byte_length"] = max(0, end_byte - heading["content_byte_start"])
    heading["content_line_count"] = max(0, end_line - heading["line"] - 1)


def iter_lines_with_offsets(
    handle: BinaryIO,
    start_offset: int,
    start_line: int,
) -> Iterable[tuple[int, int, bytes]]:
    offset = start_offset
    line_number = start_line
    handle.seek(start_offset)
    while True:
        line_start = offset
        line = handle.readline()
        if not line:
            break

        offset += len(line)
        yield line_number, line_start, line
        line_number += 1


def index_markdown(path: Path, encoding: str = "utf-8") -> dict[str, Any]:
    roots: list[dict[str, Any]] = []
    stack: list[dict[str, Any]] = []
    flat_headings: list[dict[str, Any]] = []
    in_fence = False
    fence_marker = b""
    file_size = path.stat().st_size

    with path.open("rb") as handle:
        frontmatter, body_offset, body_line = parse_frontmatter(handle, encoding)
        last_line_number = body_line - 1

        for line_number, line_start, raw_line in iter_lines_with_offsets(handle, body_offset, body_line):
            last_line_number = line_number
            line_without_eol = raw_line.rstrip(b"\r\n")

            fence_match = FENCE_RE.match(line_without_eol)
            if fence_match:
                if stack:
                    stack[-1]["content_word_count"] += count_words_in_bytes(raw_line, encoding)
                marker = fence_match.group(1)
                if not in_fence:
                    in_fence = True
                    fence_marker = marker[:1]
                elif marker.startswith(fence_marker * 3):
                    in_fence = False
                    fence_marker = b""
                continue

            if in_fence:
                if stack:
                    stack[-1]["content_word_count"] += count_words_in_bytes(raw_line, encoding)
                continue

            match = HEADING_RE.match(line_without_eol)
            if not match:
                if stack:
                    stack[-1]["content_word_count"] += count_words_in_bytes(raw_line, encoding)
                continue

            level = len(match.group(1))
            title = strip_heading_text(match.group(2), encoding)
            heading = make_heading(
                level=level,
                title=title,
                line_number=line_number,
                heading_byte_start=line_start,
                content_byte_start=line_start + len(raw_line),
            )

            if stack:
                close_heading(stack[-1], line_start, line_number)

            while stack and stack[-1]["level"] >= level:
                stack.pop()

            if stack:
                stack[-1]["children"].append(heading)
            else:
                roots.append(heading)

            stack.append(heading)
            flat_headings.append(heading)

    eof_line = last_line_number + 1 if file_size else 1

    if stack:
        close_heading(stack[-1], file_size, eof_line)

    return {
        "source_file": path.name,
        "source_size_bytes": file_size,
        "encoding": encoding,
        "frontmatter": frontmatter,
        "heading_count": len(flat_headings),
        "headings": roots,
    }


def find_markdown_files(paths: Iterable[Path], recursive: bool) -> Iterable[Path]:
    for path in paths:
        if path.is_dir():
            patterns = ("**/*.md", "**/*.markdown") if recursive else ("*.md", "*.markdown")
            for pattern in patterns:
                yield from sorted(path.glob(pattern))
        elif path.suffix.lower() in {".md", ".markdown"}:
            yield path


def output_path_for(source: Path, output_dir: Path | None) -> Path:
    filename = f"{source.name}.index.json"
    if output_dir is None:
        return source.with_name(filename)
    return output_dir / filename


def write_index(source: Path, output_dir: Path | None, encoding: str) -> Path:
    index = index_markdown(source, encoding)
    output_path = output_path_for(source, output_dir)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(index, ensure_ascii=False, indent=2, default=json_default) + "\n",
        encoding="utf-8",
    )
    return output_path


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Index Markdown front matter and ATX heading hierarchies."
    )
    parser.add_argument("paths", nargs="+", type=Path, help="Markdown files or directories")
    parser.add_argument("-o", "--output-dir", type=Path, help="directory for JSON indexes")
    parser.add_argument("-r", "--recursive", action="store_true", help="recursively find *.md files")
    parser.add_argument("--encoding", default="utf-8", help="Markdown file encoding")
    args = parser.parse_args()

    files = list(find_markdown_files(args.paths, args.recursive))
    if not files:
        raise SystemExit("No Markdown files found.")

    for source in files:
        output_path = write_index(source, args.output_dir, args.encoding)
        print(f"{source} -> {output_path}")


if __name__ == "__main__":
    main()