Robert Helewka
224541a4ce
"svg" was in both IMAGE_EXTENSIONS and PYMUPDF_EXTENSIONS, and the image
check ran first, so every SVG reached PIL.Image.open — which cannot decode
vector XML. Ingest raised UnidentifiedImageError, logged "Failed to read
image file_type=svg", and counted documents_parsed_total{status="error"}.
SVG ingest has never worked.
Rasterize to PNG instead, via PyMuPDF (already a dependency). This also
unblocks the vision stage downstream: it sends images to a vision LLM as a
data: URI, which cannot carry image/svg+xml either, so the description and
OCR fields were unreachable for SVG regardless of the parse fix.
Route SVG explicitly before both extension sets. Falling through to
PYMUPDF_EXTENSIONS would parse but not split, and a multi-page Write note
rendered as one document is either blank (no root width/height, so MuPDF
falls back to US-Letter and emits the top-left corner) or an illegible tall
strip (root sized across every stacked page). Splitting per write-page
element is correct for both formats, and since Write never rewrites
existing files both persist indefinitely.
svg_raster is a deliberate twin of daedalus's extraction/svg.py — the repos
share no common package, so the duplication is noted in both docstrings and
fixes belong in both.
lxml is a new dependency: recover=True is needed for the real files that
aren't well-formed XML, and stdlib ElementTree has no equivalent.
425 lines
15 KiB
Python
425 lines
15 KiB
Python
"""
|
|
Universal document parsing service using PyMuPDF.
|
|
|
|
Handles text extraction and image extraction for all supported formats:
|
|
PDF, EPUB, DOCX, PPTX, XLSX, XPS, MOBI, FB2, CBZ, TXT, HTML, and images.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
from dataclasses import dataclass, field
|
|
|
|
import fitz # PyMuPDF
|
|
|
|
from library.metrics import (
|
|
DOCUMENT_PARSE_DURATION,
|
|
DOCUMENTS_PARSED_TOTAL,
|
|
IMAGES_EXTRACTED_TOTAL,
|
|
)
|
|
|
|
from .text_utils import remove_excessive_whitespace, sanitize_text
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# File extensions supported by PyMuPDF
|
|
PYMUPDF_EXTENSIONS = {
|
|
"pdf", "epub", "xps", "mobi", "fb2", "cbz", "svg",
|
|
"docx", "pptx", "xlsx", "hwpx",
|
|
}
|
|
|
|
# Plain text extensions — read directly, no PyMuPDF needed
|
|
PLAINTEXT_EXTENSIONS = {"txt", "md", "csv", "tsv", "log", "json", "yaml", "yml", "xml"}
|
|
|
|
# Image extensions — store as Image nodes directly.
|
|
# SVG is deliberately absent: it is vector XML that Pillow cannot decode and
|
|
# that the vision stage cannot send as a data URI, so it gets rasterized by
|
|
# _parse_svg_file instead.
|
|
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "bmp", "tiff", "tif", "webp"}
|
|
|
|
# Minimum image dimensions to extract (skip tiny icons/bullets)
|
|
MIN_IMAGE_WIDTH = 50
|
|
MIN_IMAGE_HEIGHT = 50
|
|
|
|
|
|
@dataclass
|
|
class TextBlock:
|
|
"""A block of extracted text with page/section context."""
|
|
|
|
text: str
|
|
page: int
|
|
metadata: dict = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class ExtractedImage:
|
|
"""An image extracted from a document."""
|
|
|
|
data: bytes
|
|
ext: str
|
|
width: int
|
|
height: int
|
|
source_page: int
|
|
source_index: int
|
|
|
|
|
|
@dataclass
|
|
class ParseResult:
|
|
"""Result of parsing a document: text blocks + images + metadata."""
|
|
|
|
text_blocks: list[TextBlock] = field(default_factory=list)
|
|
images: list[ExtractedImage] = field(default_factory=list)
|
|
metadata: dict = field(default_factory=dict)
|
|
file_type: str = ""
|
|
|
|
|
|
class DocumentParser:
|
|
"""
|
|
Universal document parser using PyMuPDF.
|
|
|
|
Extracts text and images from all supported document formats through
|
|
a single unified interface.
|
|
"""
|
|
|
|
def parse(self, file_path: str, file_type: str) -> ParseResult:
|
|
"""
|
|
Parse a document and extract text blocks and images.
|
|
|
|
:param file_path: Path to the document file.
|
|
:param file_type: File extension (without dot), e.g. 'pdf', 'epub'.
|
|
:returns: ParseResult with text blocks, images, and metadata.
|
|
:raises ValueError: If the file type is not supported.
|
|
"""
|
|
file_type = file_type.lower().lstrip(".")
|
|
|
|
logger.info(
|
|
"Parsing document file_type=%s path=%s",
|
|
file_type,
|
|
os.path.basename(file_path),
|
|
)
|
|
|
|
if file_type in PLAINTEXT_EXTENSIONS:
|
|
return self._parse_plaintext(file_path, file_type)
|
|
|
|
# Checked before PYMUPDF_EXTENSIONS: PyMuPDF can open an SVG, but
|
|
# rendering a multi-page Write note as one document yields a blank or
|
|
# illegible image (see svg_raster), so it needs page-aware handling.
|
|
if file_type == "svg":
|
|
return self._parse_svg_file(file_path, file_type)
|
|
|
|
if file_type in IMAGE_EXTENSIONS:
|
|
return self._parse_image_file(file_path, file_type)
|
|
|
|
if file_type in PYMUPDF_EXTENSIONS:
|
|
return self._parse_with_pymupdf(file_path, file_type)
|
|
|
|
# HTML can be handled by PyMuPDF or direct read
|
|
if file_type in ("html", "htm"):
|
|
return self._parse_with_pymupdf(file_path, file_type)
|
|
|
|
raise ValueError(
|
|
f"Unsupported file type '{file_type}'. "
|
|
f"Supported: {sorted(PYMUPDF_EXTENSIONS | PLAINTEXT_EXTENSIONS | IMAGE_EXTENSIONS)}"
|
|
)
|
|
|
|
def parse_bytes(self, data: bytes, file_type: str, filename: str = "") -> ParseResult:
|
|
"""
|
|
Parse document from bytes (e.g. from S3 download).
|
|
|
|
:param data: Raw file bytes.
|
|
:param file_type: File extension (without dot).
|
|
:param filename: Optional original filename for logging.
|
|
:returns: ParseResult.
|
|
"""
|
|
suffix = f".{file_type.lower().lstrip('.')}"
|
|
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
|
tmp.write(data)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
logger.debug(
|
|
"Parsing from bytes file_type=%s size=%d filename=%s",
|
|
file_type,
|
|
len(data),
|
|
filename,
|
|
)
|
|
return self.parse(tmp_path, file_type)
|
|
finally:
|
|
os.unlink(tmp_path)
|
|
|
|
def _parse_with_pymupdf(self, file_path: str, file_type: str) -> ParseResult:
|
|
"""
|
|
Parse a document using PyMuPDF for text and image extraction.
|
|
|
|
:param file_path: Path to the document.
|
|
:param file_type: Normalized file extension.
|
|
:returns: ParseResult.
|
|
"""
|
|
with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
|
|
try:
|
|
doc = fitz.open(file_path)
|
|
except Exception as exc:
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
|
|
logger.error("Failed to open document file_type=%s: %s", file_type, exc)
|
|
raise
|
|
|
|
text_blocks: list[TextBlock] = []
|
|
images: list[ExtractedImage] = []
|
|
image_global_index = 0
|
|
|
|
for page_num in range(len(doc)):
|
|
page = doc[page_num]
|
|
|
|
# --- Text extraction ---
|
|
try:
|
|
text = page.get_text("text")
|
|
if text and text.strip():
|
|
cleaned = sanitize_text(text, log_changes=False)
|
|
cleaned = remove_excessive_whitespace(cleaned)
|
|
if cleaned.strip():
|
|
text_blocks.append(
|
|
TextBlock(text=cleaned, page=page_num)
|
|
)
|
|
logger.debug(
|
|
"Extracted text page=%d chars=%d",
|
|
page_num,
|
|
len(cleaned),
|
|
)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"Text extraction failed page=%d: %s, continuing",
|
|
page_num,
|
|
exc,
|
|
)
|
|
|
|
# --- Image extraction ---
|
|
try:
|
|
for img_info in page.get_images(full=True):
|
|
xref = img_info[0]
|
|
try:
|
|
img_data = doc.extract_image(xref)
|
|
if not img_data or not img_data.get("image"):
|
|
continue
|
|
|
|
width = img_data.get("width", 0)
|
|
height = img_data.get("height", 0)
|
|
|
|
# Skip tiny images (icons, bullets, etc.)
|
|
if width < MIN_IMAGE_WIDTH or height < MIN_IMAGE_HEIGHT:
|
|
logger.debug(
|
|
"Skipping small image page=%d xref=%d size=%dx%d",
|
|
page_num,
|
|
xref,
|
|
width,
|
|
height,
|
|
)
|
|
continue
|
|
|
|
images.append(
|
|
ExtractedImage(
|
|
data=img_data["image"],
|
|
ext=img_data.get("ext", "png"),
|
|
width=width,
|
|
height=height,
|
|
source_page=page_num,
|
|
source_index=image_global_index,
|
|
)
|
|
)
|
|
image_global_index += 1
|
|
|
|
logger.debug(
|
|
"Extracted image page=%d format=%s size=%dx%d bytes=%d",
|
|
page_num,
|
|
img_data.get("ext", "?"),
|
|
width,
|
|
height,
|
|
len(img_data["image"]),
|
|
)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"Image extraction failed page=%d xref=%d: %s",
|
|
page_num,
|
|
xref,
|
|
exc,
|
|
)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"Image listing failed page=%d: %s, continuing",
|
|
page_num,
|
|
exc,
|
|
)
|
|
|
|
# Collect document metadata
|
|
meta = doc.metadata or {}
|
|
result_meta = {
|
|
"page_count": len(doc),
|
|
"title": meta.get("title", ""),
|
|
"author": meta.get("author", ""),
|
|
"subject": meta.get("subject", ""),
|
|
"creator": meta.get("creator", ""),
|
|
}
|
|
|
|
doc.close()
|
|
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
|
|
IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(len(images))
|
|
|
|
logger.info(
|
|
"Parsed document file_type=%s pages=%d text_blocks=%d images=%d",
|
|
file_type,
|
|
result_meta["page_count"],
|
|
len(text_blocks),
|
|
len(images),
|
|
)
|
|
|
|
return ParseResult(
|
|
text_blocks=text_blocks,
|
|
images=images,
|
|
metadata=result_meta,
|
|
file_type=file_type,
|
|
)
|
|
|
|
def _parse_plaintext(self, file_path: str, file_type: str) -> ParseResult:
|
|
"""
|
|
Parse a plain text file by direct read.
|
|
|
|
:param file_path: Path to the text file.
|
|
:param file_type: Normalized file extension.
|
|
:returns: ParseResult.
|
|
"""
|
|
with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
|
|
try:
|
|
with open(file_path, "r", encoding="utf-8", errors="replace") as f:
|
|
content = f.read()
|
|
except Exception as exc:
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
|
|
logger.error("Failed to read text file file_type=%s: %s", file_type, exc)
|
|
raise
|
|
|
|
cleaned = sanitize_text(content, log_changes=True)
|
|
cleaned = remove_excessive_whitespace(cleaned)
|
|
|
|
text_blocks = []
|
|
if cleaned.strip():
|
|
text_blocks.append(TextBlock(text=cleaned, page=0))
|
|
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
|
|
|
|
logger.info(
|
|
"Parsed plaintext file_type=%s chars=%d",
|
|
file_type,
|
|
len(cleaned),
|
|
)
|
|
|
|
return ParseResult(
|
|
text_blocks=text_blocks,
|
|
images=[],
|
|
metadata={"page_count": 1},
|
|
file_type=file_type,
|
|
)
|
|
|
|
def _parse_svg_file(self, file_path: str, file_type: str) -> ParseResult:
|
|
"""
|
|
Rasterize an SVG into one ExtractedImage per page.
|
|
|
|
SVG is vector XML: Pillow cannot decode it and the vision stage cannot
|
|
put it in a data URI, so it is rendered to PNG here. Multi-page Write
|
|
notes become one image per page so each page reaches vision/OCR at a
|
|
legible size.
|
|
|
|
:param file_path: Path to the SVG file.
|
|
:param file_type: Normalized file extension ("svg").
|
|
:returns: ParseResult with one image per rendered page.
|
|
"""
|
|
with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
|
|
try:
|
|
from library.services.svg_raster import render_svg_pages
|
|
|
|
with open(file_path, "rb") as f:
|
|
data = f.read()
|
|
|
|
pages = render_svg_pages(data)
|
|
except Exception as exc:
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
|
|
logger.error("Failed to rasterize SVG file_type=%s: %s", file_type, exc)
|
|
raise
|
|
|
|
images = [
|
|
ExtractedImage(
|
|
data=png,
|
|
ext="png",
|
|
width=width,
|
|
height=height,
|
|
source_page=index,
|
|
source_index=0,
|
|
)
|
|
for index, (png, width, height) in enumerate(pages)
|
|
]
|
|
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
|
|
IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(len(images))
|
|
|
|
logger.info(
|
|
"Parsed SVG file_type=%s pages=%d bytes=%d",
|
|
file_type,
|
|
len(images),
|
|
len(data),
|
|
)
|
|
|
|
return ParseResult(
|
|
text_blocks=[],
|
|
images=images,
|
|
metadata={"page_count": len(images)},
|
|
file_type=file_type,
|
|
)
|
|
|
|
def _parse_image_file(self, file_path: str, file_type: str) -> ParseResult:
|
|
"""
|
|
Handle a standalone image file — store as a single ExtractedImage.
|
|
|
|
:param file_path: Path to the image file.
|
|
:param file_type: Normalized file extension.
|
|
:returns: ParseResult with one image and no text.
|
|
"""
|
|
with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
|
|
try:
|
|
from PIL import Image as PILImage
|
|
|
|
with open(file_path, "rb") as f:
|
|
data = f.read()
|
|
|
|
img = PILImage.open(file_path)
|
|
width, height = img.size
|
|
img.close()
|
|
except Exception as exc:
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
|
|
logger.error("Failed to read image file_type=%s: %s", file_type, exc)
|
|
raise
|
|
|
|
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
|
|
IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(1)
|
|
|
|
logger.info(
|
|
"Parsed image file file_type=%s size=%dx%d bytes=%d",
|
|
file_type,
|
|
width,
|
|
height,
|
|
len(data),
|
|
)
|
|
|
|
return ParseResult(
|
|
text_blocks=[],
|
|
images=[
|
|
ExtractedImage(
|
|
data=data,
|
|
ext=file_type,
|
|
width=width,
|
|
height=height,
|
|
source_page=0,
|
|
source_index=0,
|
|
)
|
|
],
|
|
metadata={"page_count": 0, "width": width, "height": height},
|
|
file_type=file_type,
|
|
)
|