mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
4e6a13576d
Addresses items 9-13 and a/b/c/f from the max-effort review of PR #272. - pdf.py: image colorspace check was `pix.n > 4`, which treats CMYK-without- alpha (n==4, same as RGBA) as not needing RGB conversion; pix.save() as .png then raises "unsupported colorspace", silently dropped by the surrounding except. Fixed to `pix.n - pix.alpha >= 4` (correctly converts CMYK, leaves RGBA untouched). - pipeline.py: detect_strategy([]) (an empty/whitespace-only source file) returned "content_based", routing into the PDF-oriented TOC-detection pipeline -- wasting a real LLM call before raising IndexingError. Empty node lists now route to level_based, whose build_tree_from_levels([]) returns an empty structure instantly with zero LLM calls. - page_index.py (shim): pageindex/__init__.py binds the canonical `page_index` function as the package attribute, but this file is ALSO a real submodule of the same name -- importing it anywhere (import machinery, unconditional) overwrites that attribute with the module object, breaking `from pageindex import page_index; page_index(x)` for the rest of the process. Made the shim module itself callable (delegates to the real function via a ModuleType subclass), so whichever object ends up in that slot is callable regardless of import order. - storage/sqlite.py: create_collection let a raw sqlite3.IntegrityError escape on a duplicate name (new CollectionAlreadyExistsError); the collections table's CHECK constraint only validated the name's first character (GLOB '*' is a wildcard, not a regex quantifier over the preceding class) -- fixed to validate the whole string, and SQLiteStorage now also validates in Python (it's a public StorageEngine usable directly, bypassing LocalBackend's own check). - tests/test_review_fixes_2.py: two tests used a ContentNode with no `level` set, so build_index took the content_based path and made real (retried, slow, and -- with a valid key -- billable) LLM calls instead of testing the text-stripping logic they claimed to. Mocked out _content_based_pipeline. - retrieve.py: _parse_pages/_get_pdf_page_content were independent copies of the canonical parse_pages/get_pdf_page_content that had already drifted (missing the p>=1 filter and 1000-page DoS cap) -- delegate to canonical now, so the legacy pageindex.get_page_content path can't silently regress again. - parser/markdown.py: a leading UTF-8 BOM broke first-header detection (not whitespace, .strip() doesn't remove it) -- decode utf-8-sig. Only backtick fences were recognized as code blocks, so a '#'-prefixed line inside a ~~~-fenced block (valid CommonMark) was misparsed as a heading -- recognize both fence styles. - run_pageindex.py: --if-thinning wasn't migrated to the bare-flag + legacy-yes/no convention the other four --if-add-* flags got; bare usage raised an argparse error and it never went through the shared coercion. - types.py: DocumentDetail's `structure` field was inside the class's total=False body, so TypedDict rules made it optional even though every backend always populates it. Split into a required base class. Adds regression tests for all of the above. Full suite: 244 passed, 2 skipped (one pre-existing, unrelated flaky cloud-streaming test). Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
78 lines
3.1 KiB
Python
78 lines
3.1 KiB
Python
import pymupdf
|
|
import pytest
|
|
from pathlib import Path
|
|
from pageindex.parser.pdf import PdfParser
|
|
from pageindex.parser.protocol import ContentNode, ParsedDocument
|
|
|
|
TEST_PDF = Path("tests/pdfs/deepseek-r1.pdf")
|
|
|
|
def test_supported_extensions():
|
|
parser = PdfParser()
|
|
assert ".pdf" in parser.supported_extensions()
|
|
|
|
@pytest.mark.skipif(not TEST_PDF.exists(), reason="Test PDF not available")
|
|
def test_parse_returns_parsed_document():
|
|
parser = PdfParser()
|
|
result = parser.parse(str(TEST_PDF))
|
|
assert isinstance(result, ParsedDocument)
|
|
assert len(result.nodes) > 0
|
|
assert result.doc_name != ""
|
|
|
|
@pytest.mark.skipif(not TEST_PDF.exists(), reason="Test PDF not available")
|
|
def test_parse_nodes_are_flat_without_level():
|
|
parser = PdfParser()
|
|
result = parser.parse(str(TEST_PDF))
|
|
for node in result.nodes:
|
|
assert isinstance(node, ContentNode)
|
|
assert node.content is not None
|
|
assert node.tokens >= 0
|
|
assert node.index is not None
|
|
assert node.level is None
|
|
|
|
|
|
def test_cmyk_pixmap_without_alpha_is_saveable_as_png(tmp_path):
|
|
"""A CMYK image with no alpha has n==4 -- same as RGBA -- so `pix.n > 4`
|
|
wrongly skips the RGB conversion PNG needs, and pix.save() raises
|
|
'unsupported colorspace for png', silently dropping the image via the
|
|
extractor's bare except. The fix (`pix.n - pix.alpha >= 4`) must convert
|
|
CMYK (4-0=4) while leaving RGBA (4-1=3) untouched."""
|
|
cmyk = pymupdf.Pixmap(pymupdf.csCMYK, pymupdf.Rect(0, 0, 10, 10))
|
|
assert cmyk.n == 4 and cmyk.alpha == 0
|
|
assert cmyk.n - cmyk.alpha >= 4 # the fixed condition: must convert
|
|
converted = pymupdf.Pixmap(pymupdf.csRGB, cmyk)
|
|
converted.save(str(tmp_path / "cmyk.png")) # must not raise
|
|
|
|
rgba = pymupdf.Pixmap(pymupdf.Pixmap(pymupdf.csRGB, pymupdf.Rect(0, 0, 10, 10)), 1)
|
|
assert rgba.n == 4 and rgba.alpha == 1
|
|
assert not (rgba.n - rgba.alpha >= 4) # unchanged: RGBA needs no conversion
|
|
rgba.save(str(tmp_path / "rgba.png")) # already saveable as-is
|
|
|
|
|
|
def test_image_paths_are_absolute(tmp_path):
|
|
"""Image references must be absolute so they resolve regardless of cwd
|
|
(cwd-relative paths broke after the query ran from another directory)."""
|
|
import os
|
|
import pymupdf
|
|
from pageindex.parser.pdf import PdfParser
|
|
|
|
# Build a 1-page PDF with an embedded image (>= _MIN_IMAGE_SIZE).
|
|
pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 64, 64), False)
|
|
pix.clear_with(128)
|
|
png = tmp_path / "img.png"
|
|
pix.save(str(png))
|
|
|
|
doc = pymupdf.open()
|
|
page = doc.new_page()
|
|
page.insert_image(pymupdf.Rect(20, 20, 180, 180), filename=str(png))
|
|
pdf_path = tmp_path / "withimg.pdf"
|
|
doc.save(str(pdf_path))
|
|
doc.close()
|
|
|
|
images_dir = tmp_path / "out" / "images"
|
|
result = PdfParser().parse(str(pdf_path), images_dir=str(images_dir))
|
|
|
|
img_paths = [im["path"] for n in result.nodes if n.images for im in n.images]
|
|
assert img_paths, "expected at least one extracted image"
|
|
for p in img_paths:
|
|
assert os.path.isabs(p), f"image path not absolute: {p}"
|
|
assert os.path.exists(p), f"image path does not resolve: {p}"
|