tracecite 0.0.2__py3-none-any.whl

tracecite/evidence/parsers/markdown.py

@@ -0,0 +1,142 @@
+"""Markdown parser: heading-hierarchy and paragraph-aware extraction.
+
+Markdown has no physical pages, so this parser synthesises a single
+retained-extraction "page" row (``physical_page=1``) holding the complete
+document text and a structural unit list in ``layout_json``, so a
+chunker-only or normalisation-only change can rebuild chunk-input units
+without re-reading the file (mirroring the PDF parser's retained-page reuse).
+"""
+
+from __future__ import annotations
+
+import re
+
+from .base import ParsedChunkUnit, ParsedPage, ParseResult
+
+NAME = "markdown-heading"
+
+_HEADING_RE = re.compile(r"^(#{1,6})\s+(.*)$")
+DEFAULT_MAX_UNIT_CHARS = 2400
+
+
+def _split_oversized(text: str, max_chars: int) -> list[str]:
+    if len(text) <= max_chars:
+        return [text]
+    words = text.split(" ")
+    parts: list[str] = []
+    current: list[str] = []
+    current_len = 0
+    for word in words:
+        word_len = len(word) + 1
+        if current and current_len + word_len > max_chars:
+            parts.append(" ".join(current))
+            current = []
+            current_len = 0
+        current.append(word)
+        current_len += word_len
+    if current:
+        parts.append(" ".join(current))
+    return parts
+
+
+def parse(path, config: dict | None = None) -> ParseResult:
+    config = config or {}
+    max_unit_chars = int(config.get("max_unit_chars", DEFAULT_MAX_UNIT_CHARS))
+
+    text = path.read_text(encoding="utf-8")
+    lines = text.splitlines()
+
+    units: list[ParsedChunkUnit] = []
+    heading_path: list[str] = []
+    heading_occurrence: dict[str, int] = {}
+    heading_key = "root#0"
+    para_index = 0
+    para_lines: list[str] = []
+    para_start_line: int | None = None
+
+    def flush_paragraph(end_line: int) -> None:
+        nonlocal para_lines, para_start_line, para_index
+        if not para_lines:
+            return
+        body = "\n".join(para_lines).strip()
+        para_lines = []
+        start_line = para_start_line
+        para_start_line = None
+        if not body:
+            return
+        for sub_index, sub_body in enumerate(_split_oversized(body, max_unit_chars)):
+            key = f"{heading_key}:p{para_index:04d}"
+            if sub_index:
+                key = f"{key}:sub{sub_index:02d}"
+            units.append(
+                ParsedChunkUnit(
+                    text=sub_body,
+                    logical_key=key,
+                    heading_path=list(heading_path),
+                    symbol=None,
+                    content_type="body",
+                    physical_page=None,
+                    line_start=start_line,
+                    line_end=end_line,
+                    locator={},
+                )
+            )
+            para_index += 1
+
+    for line_no, raw_line in enumerate(lines, start=1):
+        stripped = raw_line.strip()
+        heading_match = _HEADING_RE.match(raw_line)
+        if heading_match:
+            flush_paragraph(line_no - 1)
+            level = len(heading_match.group(1))
+            title = heading_match.group(2).strip()
+            heading_path = heading_path[: level - 1] + [title]
+            key_base = "/".join(heading_path)
+            occurrence = heading_occurrence.get(key_base, 0)
+            heading_occurrence[key_base] = occurrence + 1
+            heading_key = f"{key_base}#{occurrence}"
+            para_index = 0
+            units.append(
+                ParsedChunkUnit(
+                    text=title,
+                    logical_key=f"{heading_key}:heading",
+                    heading_path=list(heading_path),
+                    symbol=None,
+                    content_type="heading",
+                    physical_page=None,
+                    line_start=line_no,
+                    line_end=line_no,
+                    locator={},
+                )
+            )
+            continue
+        if not stripped:
+            flush_paragraph(line_no - 1)
+            continue
+        if para_start_line is None:
+            para_start_line = line_no
+        para_lines.append(raw_line)
+    flush_paragraph(len(lines))
+
+    layout_units = [
+        {
+            "text": unit.text,
+            "logical_key": unit.logical_key,
+            "heading_path": unit.heading_path,
+            "content_type": unit.content_type,
+            "line_start": unit.line_start,
+            "line_end": unit.line_end,
+        }
+        for unit in units
+    ]
+
+    page = ParsedPage(
+        physical_page=1,
+        printed_label=None,
+        text=text,
+        extraction_method=NAME,
+        extraction_status="ok" if units else "empty",
+        section_candidates=[u.text for u in units if u.content_type == "heading"],
+        layout={"units": layout_units},
+    )
+    return ParseResult(pages=[page], units=units)

tracecite/evidence/parsers/pdf.py

@@ -0,0 +1,315 @@
+"""PyMuPDF-backed PDF parser: page-, block-, and layout-aware extraction.
+
+All PyMuPDF interaction (text extraction, page rendering, image crops) is
+kept behind this module so the database and command contracts never depend
+on the specific PDF backend.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+
+import pymupdf
+
+from .base import ParsedChunkUnit, ParsedPage, ParseResult
+
+NAME = "pdf-pymupdf"
+HEADING_MAX_CHARS = 70
+
+# A page whose extracted text is shorter than this is treated as a scanned
+# page without a text layer and goes through the OCR fallback. 20 keeps a
+# near-blank title page out of OCR while staying above the few stray
+# characters PyMuPDF can extract from a rasterised page's metadata.
+MIN_EXTRACTED_TEXT_CHARS = 20
+DEFAULT_OCR_LANG = "eng"
+OCR_DPI = 300
+TESSERACT_CMD = "tesseract"
+
+# A text block is "nearby" a figure crop when it overlaps the crop
+# horizontally (with tolerance, since captions are often narrower than the
+# image) and sits within 72pt (1 inch) above or below it: close enough for a
+# caption without pulling in body text further down the page.
+NEARBY_TEXT_MAX_VERTICAL_DISTANCE = 72.0
+NEARBY_TEXT_HORIZONTAL_TOLERANCE = 20.0
+
+
+@dataclass(frozen=True)
+class RenderedAsset:
+    physical_page: int
+    asset_type: str
+    image_bytes: bytes
+    ext: str
+    width: int
+    height: int
+    bbox: tuple[float, float, float, float] | None
+    label: str | None = None
+    caption: str | None = None
+    nearby_text: str | None = None
+
+
+def _looks_like_heading(text: str) -> bool:
+    stripped = text.strip()
+    return (
+        bool(stripped)
+        and len(stripped) <= HEADING_MAX_CHARS
+        and not stripped.endswith((".", ",", ";"))
+    )
+
+
+def _blocks_for_page(page: pymupdf.Page) -> list[dict]:
+    raw_blocks = page.get_text("blocks")
+    text_blocks = [b for b in raw_blocks if b[6] == 0 and b[4].strip()]
+    text_blocks.sort(key=lambda b: (round(b[1], 1), round(b[0], 1)))
+
+    blocks: list[dict] = []
+    heading_path: list[str] = []
+    offset = 0
+    for block in text_blocks:
+        x0, y0, x1, y1, text, _block_no, _block_type = block
+        text = text.strip()
+        is_heading = _looks_like_heading(text)
+        if is_heading:
+            heading_path = [text]
+        start = offset
+        end = start + len(text)
+        offset = end + 2  # matches the "\n\n" join used to build page text
+        blocks.append(
+            {
+                "text": text,
+                "bbox": [x0, y0, x1, y1],
+                "heading_path": list(heading_path),
+                "content_type": "heading" if is_heading else "body",
+                "start_offset": start,
+                "end_offset": end,
+            }
+        )
+    return blocks
+
+
+def _nearby_text_for_bbox(
+    blocks: list[dict],
+    bbox: tuple[float, float, float, float],
+    *,
+    max_vertical_distance: float = NEARBY_TEXT_MAX_VERTICAL_DISTANCE,
+    horizontal_tolerance: float = NEARBY_TEXT_HORIZONTAL_TOLERANCE,
+) -> str | None:
+    """Find text blocks (from the page's already-computed
+    :func:`_blocks_for_page` output) vertically close to a figure crop's
+    bbox -- typically a caption directly above or below the figure -- and
+    join their text in reading order. Returns ``None`` if nothing qualifies.
+    """
+
+    x0, _y0, x1, _y1 = bbox
+    matches: list[tuple[float, str]] = []
+    for block in blocks:
+        bx0, by0, bx1, by1 = block["bbox"]
+        horizontally_overlaps = (
+            bx0 <= x1 + horizontal_tolerance and bx1 >= x0 - horizontal_tolerance
+        )
+        if not horizontally_overlaps:
+            continue
+        if by1 < bbox[1]:
+            vertical_distance = bbox[1] - by1
+        elif by0 > bbox[3]:
+            vertical_distance = by0 - bbox[3]
+        else:
+            vertical_distance = 0.0
+        if vertical_distance <= max_vertical_distance:
+            matches.append((by0, block["text"]))
+    if not matches:
+        return None
+    matches.sort(key=lambda item: item[0])
+    return "\n".join(text for _, text in matches)
+
+
+def _pymupdf_ocr_text(page: pymupdf.Page, *, lang: str, dpi: int) -> str:
+    """First OCR fallback stage: PyMuPDF's own built-in OCR support."""
+
+    try:
+        textpage = page.get_textpage_ocr(
+            flags=pymupdf.TEXTFLAGS_TEXT | pymupdf.TEXT_PRESERVE_WHITESPACE,
+            language=lang,
+            dpi=dpi,
+            full=True,
+        )
+        return page.get_text("text", textpage=textpage).strip()
+    except (RuntimeError, OSError):
+        # PyMuPDF's OCR path raises if Tesseract/tessdata is not resolvable
+        # in this environment; treat that as "no text", not a hard failure,
+        # so the caller can still try the Tesseract-CLI stage.
+        return ""
+
+
+def _tesseract_cli_ocr_text(page: pymupdf.Page, *, lang: str, dpi: int) -> str:
+    """Second OCR fallback stage: run the Tesseract CLI on a rendered page image."""
+
+    pixmap = page.get_pixmap(dpi=dpi)
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        image_path = Path(tmp_dir) / "page.png"
+        pixmap.save(str(image_path))
+        try:
+            result = subprocess.run(
+                [TESSERACT_CMD, str(image_path), "stdout", "-l", lang],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+        except (OSError, subprocess.CalledProcessError):
+            return ""
+    return result.stdout.strip()
+
+
+def _ocr_fallback_text(page: pymupdf.Page, *, lang: str, dpi: int) -> tuple[str, str]:
+    """Two-stage OCR fallback for a page whose normal text extraction came
+    back effectively empty: try PyMuPDF's built-in OCR first, and only if
+    that still returns too little text, fall back to the Tesseract CLI.
+
+    Returns ``(text, extraction_method)`` so the caller can record which
+    stage (if any) actually produced usable text.
+    """
+
+    pymupdf_text = _pymupdf_ocr_text(page, lang=lang, dpi=dpi)
+    if len(pymupdf_text) >= MIN_EXTRACTED_TEXT_CHARS:
+        return pymupdf_text, f"{NAME}-ocr-pymupdf"
+
+    tesseract_text = _tesseract_cli_ocr_text(page, lang=lang, dpi=dpi)
+    if len(tesseract_text) > len(pymupdf_text):
+        return tesseract_text, f"{NAME}-ocr-tesseract-cli"
+    return pymupdf_text, f"{NAME}-ocr-pymupdf"
+
+
+def parse(path: Path, config: dict | None = None) -> ParseResult:
+    config = config or {}
+    ocr_lang = config.get("ocr_lang", DEFAULT_OCR_LANG)
+    pages: list[ParsedPage] = []
+    units: list[ParsedChunkUnit] = []
+
+    with pymupdf.open(path) as document:
+        for index, page in enumerate(document):
+            physical_page = index + 1
+            blocks = _blocks_for_page(page)
+            page_text = "\n\n".join(block["text"] for block in blocks)
+            extraction_method = NAME
+
+            if len(page_text.strip()) < MIN_EXTRACTED_TEXT_CHARS:
+                ocr_text, ocr_method = _ocr_fallback_text(
+                    page, lang=ocr_lang, dpi=OCR_DPI
+                )
+                ocr_text = ocr_text.strip()
+                if len(ocr_text) > len(page_text.strip()):
+                    page_rect = page.rect
+                    blocks = [
+                        {
+                            "text": ocr_text,
+                            "bbox": [
+                                page_rect.x0,
+                                page_rect.y0,
+                                page_rect.x1,
+                                page_rect.y1,
+                            ],
+                            "heading_path": [],
+                            "content_type": "body",
+                            "start_offset": 0,
+                            "end_offset": len(ocr_text),
+                        }
+                    ]
+                    page_text = ocr_text
+                    extraction_method = ocr_method
+
+            section_candidates = [
+                b["text"] for b in blocks if b["content_type"] == "heading"
+            ]
+
+            pages.append(
+                ParsedPage(
+                    physical_page=physical_page,
+                    printed_label=None,
+                    text=page_text,
+                    extraction_method=extraction_method,
+                    extraction_status="ok" if blocks else "empty",
+                    section_candidates=section_candidates,
+                    layout={"blocks": blocks},
+                )
+            )
+            units.extend(_block_units(physical_page, blocks))
+
+    return ParseResult(pages=pages, units=units)
+
+
+def _block_units(physical_page: int, blocks: list[dict]) -> list[ParsedChunkUnit]:
+    return [
+        ParsedChunkUnit(
+            text=block["text"],
+            logical_key=f"page{physical_page:04d}-block{block_index:03d}",
+            heading_path=block["heading_path"],
+            symbol=None,
+            content_type=block["content_type"],
+            physical_page=physical_page,
+            page_start_offset=block["start_offset"],
+            page_end_offset=block["end_offset"],
+            locator={"bbox": block["bbox"]},
+        )
+        for block_index, block in enumerate(blocks)
+    ]
+
+
+def units_from_page_layout(
+    physical_page: int, layout_json: str | None
+) -> list[ParsedChunkUnit]:
+    """Rebuild chunk-input units from retained ``pages.layout_json``.
+
+    A chunker- or normalisation-only change rechunks from this layout
+    without reopening the PDF.
+    """
+    return _block_units(
+        physical_page, json.loads(layout_json)["blocks"] if layout_json else []
+    )
+
+
+def render_page(path: Path, physical_page: int, zoom: float = 2.0) -> RenderedAsset:
+    with pymupdf.open(path) as document:
+        page = document[physical_page - 1]
+        matrix = pymupdf.Matrix(zoom, zoom)
+        pixmap = page.get_pixmap(matrix=matrix)
+        return RenderedAsset(
+            physical_page=physical_page,
+            asset_type="page-render",
+            image_bytes=pixmap.tobytes("png"),
+            ext="png",
+            width=pixmap.width,
+            height=pixmap.height,
+            bbox=None,
+        )
+
+
+def render_figure_crops(path: Path, physical_page: int) -> list[RenderedAsset]:
+    rendered: list[RenderedAsset] = []
+    with pymupdf.open(path) as document:
+        page = document[physical_page - 1]
+        blocks = _blocks_for_page(page)
+        for image_index, image in enumerate(page.get_images(full=True)):
+            xref = image[0]
+            info = document.extract_image(xref)
+            rects = page.get_image_rects(xref)
+            bbox = tuple(rects[0]) if rects else None
+            nearby_text = (
+                _nearby_text_for_bbox(blocks, bbox) if bbox is not None else None
+            )
+            rendered.append(
+                RenderedAsset(
+                    physical_page=physical_page,
+                    asset_type="figure-crop",
+                    image_bytes=info["image"],
+                    ext=info["ext"],
+                    width=info.get("width", 0),
+                    height=info.get("height", 0),
+                    bbox=bbox,
+                    label=f"page-{physical_page:03d}-image-{image_index:02d}",
+                    nearby_text=nearby_text,
+                )
+            )
+    return rendered