feat(F-09): rich parsing — formula + figure extraction
- codex/parsing/mathpix.py: pix2tex (local, CPU) primary + MathPix API optional; bbox heuristic h>15px, math-char-count>5; singleton model cache - codex/parsing/figures.py: pymupdf embedded-image extraction → PNG; caption detection via proximity + "Figure/Fig./Abbildung" prefix - codex/models.py: FormulaChunk + FigureChunk dataclasses (R-10/R-11) - codex/ingest.py: --rich flag wires formula+figure extraction post-ingest - codex/cli.py: search_app sub-typer (paper + formula subcommands), --rich flag on ingest; wiki_app from F-12 preserved intact - codex/config.py: mathpix_app_id/key, pix2tex_fallback, figures_dir - infra/schema.sql: formulas + figures tables with HNSW pgvector indexes - pyproject.toml: pymupdf>=1.24, pix2tex>=0.1.4 - tests/parsing/test_mathpix.py + test_figures.py: 31 tests (mock pix2tex + MathPix HTTP, real pymupdf on synthetic PDF) Gate: 158 passed, ruff clean, mypy clean (20 files) Requirements: R-10 R-11 R-12 R-13 R-14 → done Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Tarik Moussa
245
tests/parsing/test_figures.py
Normal file
245
tests/parsing/test_figures.py
Normal file
@@ -0,0 +1,245 @@
|
||||
"""Tests for codex.parsing.figures — figure extraction module.
|
||||
|
||||
All external dependencies (fitz, PIL) are mocked so the suite runs offline.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
def _make_mock_rect(
|
||||
x0: float = 10.0, y0: float = 10.0, x1: float = 200.0, y1: float = 150.0
|
||||
) -> MagicMock:
|
||||
"""Build a mock fitz.Rect-like object."""
|
||||
rect = MagicMock()
|
||||
rect.x0 = x0
|
||||
rect.y0 = y0
|
||||
rect.x1 = x1
|
||||
rect.y1 = y1
|
||||
return rect
|
||||
|
||||
|
||||
def _make_mock_doc(
|
||||
images: list[tuple[int, ...]], # list of (xref, ...)
|
||||
text_blocks: list[Any],
|
||||
img_rects: list[Any],
|
||||
page_count: int = 1,
|
||||
) -> tuple[MagicMock, MagicMock]:
|
||||
"""Build a mock fitz document with one page."""
|
||||
mock_pix = MagicMock()
|
||||
mock_pix.n = 3 # RGB — no conversion needed
|
||||
mock_pix.save = MagicMock()
|
||||
|
||||
mock_page = MagicMock()
|
||||
mock_page.get_images.return_value = images
|
||||
mock_page.get_text.return_value = text_blocks
|
||||
mock_page.get_image_rects.return_value = img_rects
|
||||
|
||||
mock_doc = MagicMock()
|
||||
mock_doc.__len__.return_value = page_count
|
||||
mock_doc.__getitem__ = MagicMock(return_value=mock_page)
|
||||
|
||||
return mock_doc, mock_pix
|
||||
|
||||
|
||||
class TestExtractFigures:
|
||||
"""Tests for extract_figures()."""
|
||||
|
||||
def test_no_caption_returns_empty_string(self, tmp_path: Any) -> None:
|
||||
"""Figure with no nearby caption gets empty caption string."""
|
||||
img_rect = _make_mock_rect(10, 10, 200, 150)
|
||||
text_blocks = [
|
||||
(300.0, 300.0, 400.0, 320.0, "Unrelated text here.", 0, 0),
|
||||
]
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=text_blocks,
|
||||
img_rects=[img_rect],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].caption == ""
|
||||
|
||||
def test_figure_png_is_saved(self, tmp_path: Any) -> None:
|
||||
"""PNG save is called with the correct output path."""
|
||||
img_rect = _make_mock_rect()
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=[],
|
||||
img_rects=[img_rect],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert len(results) == 1
|
||||
mock_pix.save.assert_called_once()
|
||||
saved_path: str = mock_pix.save.call_args[0][0]
|
||||
assert saved_path.endswith(".png")
|
||||
|
||||
def test_caption_detected_below_image(self, tmp_path: Any) -> None:
|
||||
"""Caption block immediately below the image is detected."""
|
||||
img_rect = _make_mock_rect(10, 10, 200, 150)
|
||||
caption_block = (10.0, 155.0, 200.0, 170.0, "Figure 1: Architecture overview.", 0, 0)
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=[caption_block],
|
||||
img_rects=[img_rect],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert results[0].caption == "Figure 1: Architecture overview."
|
||||
|
||||
def test_caption_detected_above_image(self, tmp_path: Any) -> None:
|
||||
"""Caption block immediately above the image is also detected."""
|
||||
img_rect = _make_mock_rect(10, 80, 200, 200)
|
||||
caption_block = (10.0, 50.0, 200.0, 75.0, "Fig. 2: Loss curve.", 0, 0)
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=[caption_block],
|
||||
img_rects=[img_rect],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert results[0].caption == "Fig. 2: Loss curve."
|
||||
|
||||
def test_german_caption_prefix(self, tmp_path: Any) -> None:
|
||||
"""'Abbildung' prefix is recognized as a valid German caption."""
|
||||
img_rect = _make_mock_rect(10, 10, 200, 150)
|
||||
caption_block = (10.0, 155.0, 200.0, 170.0, "Abbildung 3: Ergebnisse.", 0, 0)
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=[caption_block],
|
||||
img_rects=[img_rect],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert "Abbildung" in results[0].caption
|
||||
|
||||
def test_unrelated_text_not_used_as_caption(self, tmp_path: Any) -> None:
|
||||
"""Text blocks not starting with a caption prefix are ignored."""
|
||||
img_rect = _make_mock_rect(10, 10, 200, 150)
|
||||
non_caption = (10.0, 155.0, 200.0, 170.0, "This is a random sentence.", 0, 0)
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=[non_caption],
|
||||
img_rects=[img_rect],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert results[0].caption == ""
|
||||
|
||||
def test_no_images_returns_empty_list(self, tmp_path: Any) -> None:
|
||||
"""PDF with no embedded images returns empty list."""
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[],
|
||||
text_blocks=[],
|
||||
img_rects=[],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_image_without_rect_is_skipped(self, tmp_path: Any) -> None:
|
||||
"""Images for which get_image_rects returns empty list are skipped."""
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=[],
|
||||
img_rects=[], # no rect
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("paper.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_output_dir_created_if_absent(self, tmp_path: Any) -> None:
|
||||
"""output_dir is created if it does not exist."""
|
||||
new_dir = tmp_path / "new_subdir" / "figures"
|
||||
assert not new_dir.exists()
|
||||
|
||||
mock_doc, mock_pix = _make_mock_doc(images=[], text_blocks=[], img_rects=[])
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
extract_figures("paper.pdf", output_dir=str(new_dir))
|
||||
|
||||
assert new_dir.exists()
|
||||
|
||||
def test_paper_id_from_pdf_stem(self, tmp_path: Any) -> None:
|
||||
"""paper_id in FigureChunk is the stem of the PDF filename."""
|
||||
img_rect = _make_mock_rect()
|
||||
mock_doc, mock_pix = _make_mock_doc(
|
||||
images=[(5, 0, 0, 0, 0, 0, 0, 0, 8, "png")],
|
||||
text_blocks=[],
|
||||
img_rects=[img_rect],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("fitz.open", return_value=mock_doc),
|
||||
patch("fitz.Pixmap", return_value=mock_pix),
|
||||
):
|
||||
from codex.parsing.figures import extract_figures
|
||||
|
||||
results = extract_figures("/tmp/2301.07041.pdf", output_dir=str(tmp_path))
|
||||
|
||||
assert results[0].paper_id == "2301.07041"
|
||||
Reference in New Issue
Block a user