extract_pdf_annotations.py

"""Extract PDF text-markup annotations into JSON.

Usage: python3.14 extract_pdf_annotations.py INPUT.pdf [-o OUTPUT.json] [--pretty]
Requires the ``PyMuPDF`` package, imported as ``fitz``.
"""

from __future__ import annotations

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

import fitz  # PyMuPDF

TEXT_MARKUP_TYPES = {
    fitz.PDF_ANNOT_HIGHLIGHT: "highlight", # type: ignore
    fitz.PDF_ANNOT_UNDERLINE: "underline", # type: ignore
    fitz.PDF_ANNOT_SQUIGGLY: "squiggly", # type: ignore
    fitz.PDF_ANNOT_STRIKE_OUT: "strikeout", # type: ignore
}


def get_annotation_popup_text(annot: fitz.Annot) -> str | None:
    """Return the annotation's note/comment text, when present."""
    info = annot.info or {}
    content = info.get("content")
    if content is not None:
        content = content.strip()
    return content or None


def _pair_to_point(value: Any) -> tuple[float, float] | None:
    if isinstance(value, fitz.Point):
        return float(value.x), float(value.y)
    if isinstance(value, (list, tuple)) and len(value) >= 2:
        x, y = value[0], value[1]
        if isinstance(x, (int, float)) and isinstance(y, (int, float)):
            return float(x), float(y)
    return None


def normalize_quads(vertices: Any) -> list[list[tuple[float, float]]]:
    if not vertices:
        return []

    if isinstance(vertices, fitz.Quad):
        return [[
            (float(vertices.ul.x), float(vertices.ul.y)),
            (float(vertices.ur.x), float(vertices.ur.y)),
            (float(vertices.ll.x), float(vertices.ll.y)),
            (float(vertices.lr.x), float(vertices.lr.y)),
        ]]

    if not isinstance(vertices, Iterable) or isinstance(vertices, (str, bytes)):
        return []

    items = list(vertices)
    if not items:
        return []

    if all(isinstance(item, (int, float)) for item in items):
        points = [
            (float(items[index]), float(items[index + 1]))
            for index in range(0, len(items) - 1, 2)
        ]
        return [
            points[index:index + 4]
            for index in range(0, len(points), 4)
            if len(points[index:index + 4]) == 4
        ]

    normalized_points = [_pair_to_point(item) for item in items]
    if all(point is not None for point in normalized_points):
        points = [point for point in normalized_points if point is not None]
        return [
            points[index:index + 4]
            for index in range(0, len(points), 4)
            if len(points[index:index + 4]) == 4
        ]

    quads: list[list[tuple[float, float]]] = []
    for item in items:
        if isinstance(item, fitz.Quad):
            quads.append([
                (float(item.ul.x), float(item.ul.y)),
                (float(item.ur.x), float(item.ur.y)),
                (float(item.ll.x), float(item.ll.y)),
                (float(item.lr.x), float(item.lr.y)),
            ])
            continue

        if isinstance(item, Iterable) and not isinstance(item, (str, bytes)):
            quad_points = [_pair_to_point(point) for point in item]
            if len(quad_points) >= 4 and all(point is not None for point in quad_points[:4]):
                quads.append([point for point in quad_points[:4] if point is not None])

    return quads


def extract_text_from_quads(page: fitz.Page, quads: list[list[tuple[float, float]]]) -> str:
    """
    Extract text from markup quad points.
    """
    parts = []
    for quad in quads:
        # Convert the four quad points to a rectangle.
        if len(quad) < 4:
            continue

        xs = [point[0] for point in quad]
        ys = [point[1] for point in quad]
        rect = fitz.Rect(min(xs), min(ys), max(xs), max(ys))

        # Extract text from this area.
        text = page.get_text("text", clip=rect)
        if text:
            parts.append(text.strip()) # type: ignore

    # Remove consecutive duplicates.
    deduped = []
    for part in parts:
        if not deduped or deduped[-1] != part:
            deduped.append(part)

    return " ".join(deduped).strip()


def extract_annotations(pdf_path: str) -> dict[str, Any]:
    doc = fitz.open(pdf_path)

    result: dict[str, Any] = {
        "source_file": str(Path(pdf_path).resolve()),
        "page_count": doc.page_count,
        "annotations": [],
    }

    total_index = 0

    for page_number in range(doc.page_count):
        page = doc[page_number]
        annot = page.first_annot
        page_annot_index = 0

        while annot:
            page_annot_index += 1
            total_index += 1

            annot_type_code = annot.type[0]
            annot_type_name = TEXT_MARKUP_TYPES.get(annot_type_code, "unknown")
            comment_text = get_annotation_popup_text(annot)
            vertices = getattr(annot, "vertices", None)
            quads: list[list[tuple[float, float]]] = []

            # Text-markup annotations expose their selected text as quad points.
            if annot_type_code in TEXT_MARKUP_TYPES:
                quads = normalize_quads(vertices)
                selected_text = extract_text_from_quads(page, quads)
            else:
                selected_text = None

            entry: dict[str, Any] = {
                "annotation_id": f"p{page_number + 1}_a{page_annot_index}",
                "global_index": total_index,
                "page": page_number + 1,
                "annotation_index": page_annot_index,
                "type": annot_type_name,
                "selected_text": selected_text,
                "comment": comment_text,
            }

            result["annotations"].append(entry)
            annot = annot.next

    doc.close()
    return result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract PDF annotations to JSON."
    )
    parser.add_argument("input_pdf", help="input PDF path")
    parser.add_argument(
        "-o",
        "--output",
        help="output JSON path (default: next to the PDF)",
        default=None,
    )
    parser.add_argument(
        "--pretty",
        action="store_true",
        help="write indented, human-readable JSON",
    )
    args = parser.parse_args()

    input_pdf = Path(args.input_pdf)
    if not input_pdf.exists():
        raise FileNotFoundError(f"File not found: {input_pdf}")

    output_path = (
        Path(args.output)
        if args.output
        else input_pdf.with_suffix(".annotations.json")
    )

    data = extract_annotations(str(input_pdf))

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(
            data,
            f,
            ensure_ascii=False,
            indent=2 if args.pretty else None,
        )

    print(f"Done: {output_path}")


if __name__ == "__main__":
    main()