diff --git a/paperforge/worker/ocr_emit.py b/paperforge/worker/ocr_emit.py new file mode 100644 index 00000000..544f1e47 --- /dev/null +++ b/paperforge/worker/ocr_emit.py @@ -0,0 +1,52 @@ +from __future__ import annotations + + +def emit_page_markdown( + page_index: int, + ordered_spine: list, + attachments: list, + page_width: int = 0, + page_height: int = 0, +) -> str: + lines: list[str] = [f""] + attachment_by_anchor: dict[str, list] = {} + + for att in attachments: + anchor = getattr(att, "anchor_node_id", "") if hasattr(att, "anchor_node_id") else att.get("anchor_node_id", "") + if anchor: + attachment_by_anchor.setdefault(anchor, []).append(att) + + inserted_ids: set = set() + + for node in ordered_spine: + node_id = ( + getattr(node, "node_id", "") + if hasattr(node, "node_id") + else node.get("node_id", "") + ) + node_type = ( + getattr(node, "node_type", "paragraph") + if hasattr(node, "node_type") + else node.get("node_type", "paragraph") + ) + text = ( + getattr(node, "text", "") + if hasattr(node, "text") + else node.get("text", "") + ) + + if node_id in attachment_by_anchor: + for att in attachment_by_anchor[node_id]: + att_id = getattr(att, "attachment_id", id(att)) if hasattr(att, "attachment_id") else att.get("attachment_id", id(att)) + if att_id not in inserted_ids: + inserted_ids.add(att_id) + kind = getattr(att, "kind", "figure") if hasattr(att, "kind") else att.get("kind", "figure") + lines.append(f"[{kind.upper()} attached: {att_id}]") + + if node_type in ("section_heading", "subsection_heading"): + lines.append(f"### {text}") + else: + if text: + lines.append(text) + + return "\n".join(lines) diff --git a/tests/test_ocr_emission_regressions.py b/tests/test_ocr_emission_regressions.py new file mode 100644 index 00000000..7dd0e6c7 --- /dev/null +++ b/tests/test_ocr_emission_regressions.py @@ -0,0 +1,66 @@ +from __future__ import annotations + + +def test_emission_keeps_headings_in_markdown() -> None: + from paperforge.worker.ocr_emit import emit_page_markdown + + class Node: + def __init__(self, node_id, node_type, text=""): + self.node_id = node_id + self.node_type = node_type + self.text = text + + spine = [ + Node("h1", "section_heading", "5 In vitro PEMF studies"), + Node("p1", "paragraph", "Over the past few decades..."), + ] + + result = emit_page_markdown(7, spine, []) + + assert "### 5 In vitro PEMF studies" in result + assert "" in result + assert "Over the past few decades..." in result + + +def test_emission_preserves_page_marker() -> None: + from paperforge.worker.ocr_emit import emit_page_markdown + + result = emit_page_markdown(14, [], []) + + assert result.startswith("") + + +def test_emission_handles_empty_input() -> None: + from paperforge.worker.ocr_emit import emit_page_markdown + + result = emit_page_markdown(1, [], []) + + assert "" in result + + +def test_emission_attachments_referenced_near_anchor() -> None: + from paperforge.worker.ocr_emit import emit_page_markdown + + class Node: + def __init__(self, node_id, node_type, text=""): + self.node_id = node_id + self.node_type = node_type + self.text = text + + class Att: + def __init__(self, attachment_id, kind, anchor_node_id): + self.attachment_id = attachment_id + self.kind = kind + self.anchor_node_id = anchor_node_id + + spine = [ + Node("sec-6-1", "section_heading", "6.1 Bone"), + ] + attachments = [ + Att("a1", "figure", "sec-6-1"), + ] + + result = emit_page_markdown(14, spine, attachments) + + assert "### 6.1 Bone" in result + assert "[FIGURE attached: a1]" in result