"""Runnable check for source classification and raster loading. Builds its own PDFs so it needs no corpus files (scores are copyrighted and gitignored). Run with `python tests/test_pdf.py`. """ from __future__ import annotations import sys from pathlib import Path import numpy as np import pymupdf sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from noteman_slicer.pdf import SourceType, open_source, page_raster # noqa: E402 A4 = pymupdf.paper_rect("a4") def _vector_pdf(path: Path, pages: int = 2) -> None: doc = pymupdf.open() for _ in range(pages): page = doc.new_page(width=A4.width, height=A4.height) page.draw_line((50, 100), (A4.width - 50, 100)) page.insert_text((50, 150), "notation", fontsize=24) doc.save(path) def _scan_pdf(path: Path, pages: int = 2, w: int = 1653, h: int = 2332) -> None: """Each page is one full-page grayscale image — what a real scan looks like.""" art = np.full((h, w), 255, np.uint8) art[500:505, 100 : w - 100] = 0 # a staff line, so it isn't uniform pix = pymupdf.Pixmap(pymupdf.csGRAY, w, h, bytearray(art.tobytes()), False) doc = pymupdf.open() for _ in range(pages): page = doc.new_page(width=A4.width, height=A4.height) page.insert_image(page.rect, pixmap=pix) doc.save(path) def main() -> int: tmp = Path(__file__).with_name("_tmp") tmp.mkdir(exist_ok=True) vec, scan = tmp / "vector.pdf", tmp / "scan.pdf" _vector_pdf(vec) _scan_pdf(scan) src = open_source(vec) assert src.type is SourceType.VECTOR, src.type assert not src.overridden page = page_raster(src, 0) # Rendered at 600 DPI, so an A4 page is ~4960px wide. assert page.ndim == 2 and page.dtype == np.uint8, (page.ndim, page.dtype) assert 4900 < page.shape[1] < 5000, page.shape src.close() src = open_source(scan) assert src.type is SourceType.RASTER, src.type page = page_raster(src, 0) # Native resolution of the embedded image, NOT a 600 DPI re-render. assert page.shape == (2332, 1653), page.shape assert page.min() == 0 and page.max() == 255, (page.min(), page.max()) src.close() # An override must win over detection, and say so. src = open_source(scan, SourceType.VECTOR) assert src.type is SourceType.VECTOR and src.detected is SourceType.RASTER assert src.overridden assert page_raster(src, 0).shape[1] > 4000, "override must force a render" src.close() for f in (vec, scan): f.unlink() tmp.rmdir() print("ok") return 0 if __name__ == "__main__": sys.exit(main())