From 5838d8b49e904ac3b8413a86151a7d3498660c2f Mon Sep 17 00:00:00 2001 From: myhloli Date: Sat, 5 Sep 2026 14:38:36 +0800 Subject: [PATCH 1/2] refactor: extract Flash PDF pages once and isolate pdftext adapters --- mineru/model/flash/pdf/document.py | 327 +++++++++------------- mineru/model/flash/pdf/native_text.py | 10 +- mineru/model/flash/pdf/pdftext_adapter.py | 175 ++++++++++++ mineru/model/flash/pdf/pipeline.py | 42 ++- tests/unittest/test_pdf_document.py | 92 ++++++ 5 files changed, 424 insertions(+), 222 deletions(-) create mode 100644 mineru/model/flash/pdf/pdftext_adapter.py diff --git a/mineru/model/flash/pdf/document.py b/mineru/model/flash/pdf/document.py index b78c4a424..015c0c0c6 100644 --- a/mineru/model/flash/pdf/document.py +++ b/mineru/model/flash/pdf/document.py @@ -12,12 +12,10 @@ from io import BytesIO from typing import Any, Iterator, Literal, TypeAlias, cast -import numpy as np import pypdfium2 as pdfium import pypdfium2.raw as pdfium_c -from pdftext.pdf.chars import deduplicate_chars, get_chars -from pdftext.pdf.pages import assign_scripts, get_lines, get_spans -from pdftext.schema import Bbox, Char, Line +from pdftext.pdf.chars import get_chars +from pdftext.schema import Char, Line from PIL import Image, ImageOps from ....types import BBox, PageInfo @@ -26,6 +24,8 @@ from .._shared.image import image_to_bytes from .classify import classify from .pdfium import _pdfium_lock +from .pdftext_adapter import _char_bbox_values, _deduplicate_pdftext_chars, _ensure_legacy_chars +from .pdftext_adapter import get_lines_from_chars as get_lines_from_chars logger = logging.getLogger(__name__) @@ -47,10 +47,6 @@ PDF_IMAGE_FINGERPRINT_MAX_RAW_BYTES = 16 * 1024 * 1024 _PDF_EXTERNAL_LINK_SCHEMES = frozenset({"http", "https", "mailto", "tel"}) -try: - from pdftext.pdf.chars import PageChars -except ImportError: - PageChars = None # See: pdfium.PdfDocument.METADATA_KEYS PDFMetadataKey: TypeAlias = Literal[ @@ -122,6 +118,21 @@ class PDFImageInfo: fingerprint: str | None +@dataclass(frozen=True, slots=True) +class _PDFPageSnapshot: + """保存单次页面生命周期提取的纯 Python 数据，不持有 PDFium 子对象。""" + + page_size: tuple[float, float] + rotation: Literal[0, 90, 180, 270] + text_geometry: PDFPageTextGeometry + drawing_lines: list[PDFDrawingLine] + path_infos: list[PDFPathInfo] + image_infos: list[PDFImageInfo] + form_bboxes: list[BBox] + signature_bboxes: list[BBox] + link_annotations: list[PDFLinkAnnotation] + + @dataclass class _PathSubpath: """保存一个 PDF Path 子路径的点、直线段和闭合状态。""" @@ -400,6 +411,34 @@ def get_page_chars_with_geometry(self, page_idx: int) -> PDFPageTextGeometry: """一次读取字符、tight bbox 和字符原点，避免 Hybrid 重复打开 textpage。""" return self._get_page_text_geometry(page_idx, include_extended_geometry=True) + def _extract_native_page(self, page_idx: int) -> _PDFPageSnapshot: + """在一次加锁打开中收集 Flash 页面证据，并共享 Path 子路径解码。""" + + with self._open_page(page_idx) as page: + page_bbox = _normalize_pdf_page_bbox(page.get_bbox()) + try: + raw_rotation = int(page.get_rotation()) % 360 + except Exception: + raw_rotation = 0 + rotation = raw_rotation if raw_rotation in {0, 90, 180, 270} else 0 + drawings, paths = _extract_page_paths_and_lines(page, page_bbox, raw_rotation) + return _PDFPageSnapshot( + page_size=_drawing_page_size(page_bbox, raw_rotation), + rotation=cast(Literal[0, 90, 180, 270], rotation), + text_geometry=_extract_page_text_geometry(page, include_extended_geometry=True), + drawing_lines=drawings, + path_infos=paths, + image_infos=_extract_page_image_infos(page, page_bbox, raw_rotation), + form_bboxes=_extract_page_form_bboxes(page, page_bbox, raw_rotation), + signature_bboxes=_extract_page_signature_bboxes( + page, + page_bbox, + raw_rotation, + form_handle=getattr(self._pdf_doc, "formenv", None), + ), + link_annotations=_extract_page_link_annotations(page, self._pdf_doc.raw, page_bbox, raw_rotation), + ) + def _get_page_text_geometry( self, page_idx: int, @@ -408,38 +447,7 @@ def _get_page_text_geometry( ) -> PDFPageTextGeometry: """在同一 PDFium textpage 生命周期内物化字符及可选扩展几何。""" with self._open_page(page_idx) as page: - textpage = None - try: - textpage = page.get_textpage() - raw_page_bbox: list[float] = list(page.get_bbox()) - page_bbox = _normalize_pdf_page_bbox(tuple(raw_page_bbox)) - page_rotation: int = 0 - try: - page_rotation = page.get_rotation() - except Exception: - pass - chars = get_chars(textpage, raw_page_bbox, page_rotation) - chars = _deduplicate_pdftext_chars(chars) - chars = _ensure_legacy_chars(chars) - chars = _restore_pdfium_surrogate_pairs(chars, textpage) - chars = _deduplicate_near_identical_chars(chars) - if include_extended_geometry: - loose_bboxes, tight_bboxes, origins = _extract_page_char_extended_geometry( - textpage, - chars, - page_bbox, - page_rotation, - ) - else: - loose_bboxes, tight_bboxes, origins = {}, {}, {} - finally: - _try_close(textpage) - return PDFPageTextGeometry( - chars=chars, - tight_bboxes=tight_bboxes, - origins=origins, - loose_bboxes=loose_bboxes, - ) + return _extract_page_text_geometry(page, include_extended_geometry=include_extended_geometry) def get_page_lines(self, page_idx: int) -> list[Line]: chars = self.get_page_chars(page_idx) @@ -1104,6 +1112,8 @@ def _path_info_from_object( page_rotation: int, form_depth: int, source_index: int, + *, + subpaths: list[_PathSubpath] | None = None, ) -> PDFPathInfo | None: """将一个原始 Path 转换为页面几何；贝塞尔控制点也纳入保守 bbox。""" @@ -1117,7 +1127,7 @@ def _path_info_from_object( fill_visible, stroke_visible = _get_path_visibility(raw_obj) points = [ point - for raw_subpath in _read_raw_path_subpaths(raw_obj) + for raw_subpath in (_read_raw_path_subpaths(raw_obj) if subpaths is None else subpaths) for point in _transform_path_subpath( raw_subpath, matrix, @@ -1265,6 +1275,8 @@ def _extract_path_drawing_lines( matrix: tuple[float, float, float, float, float, float], page_bbox: BBox, page_rotation: int, + *, + subpaths: list[_PathSubpath] | None = None, ) -> list[PDFDrawingLine]: """从单个 Path 提取可见直线，坏 Path 返回空结果且不影响同页其他对象。""" fill_visible, stroke_visible = _get_path_visibility(raw_obj) @@ -1274,7 +1286,7 @@ def _extract_path_drawing_lines( page_size = _drawing_page_size(page_bbox, page_rotation) raw_stroke_width = _get_raw_stroke_width(raw_obj) if stroke_visible else 0.0 drawing_lines: list[PDFDrawingLine] = [] - for raw_subpath in _read_raw_path_subpaths(raw_obj): + for raw_subpath in _read_raw_path_subpaths(raw_obj) if subpaths is None else subpaths: subpath = _transform_path_subpath(raw_subpath, matrix, page_bbox, page_rotation) if fill_visible: filled_line = _get_thin_filled_subpath_line(subpath, page_size) @@ -1937,6 +1949,41 @@ def _extract_page_path_infos( return path_infos +def _extract_page_paths_and_lines( + page: pdfium.PdfPage, + page_bbox: BBox, + page_rotation: int, +) -> tuple[list[PDFDrawingLine], list[PDFPathInfo]]: + """一次遍历和解码 Path，分别隔离绘图线与路径信息的派生异常。""" + + drawing_lines: list[PDFDrawingLine] = [] + path_infos: list[PDFPathInfo] = [] + for source_index, (raw_obj, matrix, form_depth) in enumerate(_iter_raw_path_objects_with_depth(page)): + try: + subpaths = _read_raw_path_subpaths(raw_obj) + except Exception: + continue + try: + drawing_lines.extend(_extract_path_drawing_lines(raw_obj, matrix, page_bbox, page_rotation, subpaths=subpaths)) + except Exception: + pass + try: + info = _path_info_from_object( + raw_obj, + matrix, + page_bbox, + page_rotation, + form_depth, + source_index, + subpaths=subpaths, + ) + except Exception: + continue + if info is not None: + path_infos.append(info) + return _merge_collinear_drawing_lines(drawing_lines, _drawing_page_size(page_bbox, page_rotation)), path_infos + + def _extract_page_drawing_lines( page: pdfium.PdfPage, page_bbox: BBox, @@ -2119,47 +2166,6 @@ def _deduplicate_near_identical_chars(chars: list[Char]) -> list[Char]: return deduplicated_chars -def _is_pdftext_page_chars(chars: Any) -> bool: - """判断对象是否为 pdftext 0.7 引入的 PageChars 列式字符容器。""" - return PageChars is not None and isinstance(chars, PageChars) - - -def _deduplicate_pdftext_chars(chars: Any) -> Any: - """按当前 pdftext 返回类型调用官方去重，兼容测试或旧版本的 list 字符。""" - if _is_pdftext_page_chars(chars) or PageChars is None: - return deduplicate_chars(chars) - return chars - - -def _materialize_page_chars(chars: Any) -> list[Char]: - """将 pdftext 0.7 的 PageChars 物化为 MinerU 既有 char dict 列表。""" - boxes = chars.boxes.tolist() - rotations = chars.rotations.tolist() - font_ids = chars.font_ids.tolist() - char_indices = chars.char_indices.tolist() - - return [ - cast( - Char, - { - "bbox": Bbox([float(coord) for coord in boxes[index]]), - "char": chars.text[index], - "rotation": float(rotations[index]), - "font": chars.fonts[int(font_ids[index])], - "char_idx": int(char_indices[index]), - }, - ) - for index in range(len(chars)) - ] - - -def _ensure_legacy_chars(chars: Any) -> list[Char]: - """统一输出旧版 char dict 列表，隔离 pdftext 0.7 的返回结构变化。""" - if _is_pdftext_page_chars(chars): - return _materialize_page_chars(chars) - return cast(list[Char], chars) - - def _restore_pdfium_surrogate_pairs( chars: list[Char], textpage: pdfium.PdfTextPage, @@ -2232,121 +2238,6 @@ def _restore_pdfium_surrogate_pairs( return restored_chars -def _char_bbox_values(bbox: object) -> tuple[float, float, float, float] | None: - """将 tuple/list 或 pdftext Bbox 对象统一转换为四元组坐标。""" - if bbox is None: - return None - if isinstance(bbox, (list, tuple)) and len(bbox) == 4: - return tuple(float(value) for value in bbox) # type: ignore[return-value] - - attrs = ("x_start", "y_start", "x_end", "y_end") - if all(hasattr(bbox, attr) for attr in attrs): - return tuple(float(getattr(bbox, attr)) for attr in attrs) # type: ignore[return-value] - return None - - -def _get_single_char_text(char: Char) -> str: - """提取单个 PDF 字符文本，异常空值用替换符保证 PageChars 长度一致。""" - text = str(char.get("char", "")) - if len(text) == 1: - return text - return text[:1] or "\ufffd" - - -def _get_char_font_id( - char: Char, - fonts: list[dict[str, Any]], - font_cache: dict[tuple[Any, Any, Any, Any], int], -) -> int: - """为旧版字符 font 生成 PageChars 需要的页内 font id。""" - font = char.get("font") or {} - font_key = ( - font.get("name"), - font.get("flags"), - font.get("size"), - font.get("weight"), - ) - font_id = font_cache.get(font_key) - if font_id is None: - font_id = len(fonts) - font_cache[font_key] = font_id - fonts.append( - { - "name": font.get("name"), - "flags": font.get("flags"), - "size": font.get("size"), - "weight": font.get("weight"), - } - ) - return font_id - - -def _get_char_index(char: Char, fallback_idx: int) -> int: - """提取旧版字符索引，缺失或为空时回退到当前列表位置。""" - char_idx = char.get("char_idx") - if char_idx is None: - char_idx = fallback_idx - return int(char_idx) - - -def _legacy_chars_to_page_chars(chars: Any) -> Any: - """将旧版 char dict 列表打包回 pdftext 0.7 get_spans 所需的 PageChars。""" - if PageChars is None or _is_pdftext_page_chars(chars): - return chars - - fonts: list[dict[str, Any]] = [] - font_cache: dict[tuple[Any, Any, Any, Any], int] = {} - text_parts: list[str] = [] - codes: list[int] = [] - rotations: list[float] = [] - boxes: list[tuple[float, float, float, float]] = [] - font_ids: list[int] = [] - char_indices: list[int] = [] - - for fallback_idx, char in enumerate(cast(list[Char], chars)): - char_text = _get_single_char_text(char) - bbox_values = _char_bbox_values(char.get("bbox")) - if bbox_values is None: - bbox_values = (0.0, 0.0, 0.0, 0.0) - text_parts.append(char_text) - codes.append(ord(char_text)) - rotations.append(float(char.get("rotation") or 0.0)) - boxes.append(bbox_values) - font_ids.append(_get_char_font_id(char, fonts, font_cache)) - char_indices.append(_get_char_index(char, fallback_idx)) - - return PageChars( - "".join(text_parts), - np.array(codes, dtype=np.uint32), - np.array(rotations, dtype=np.float64), - np.array(boxes, dtype=np.float64).reshape((len(boxes), 4)), - np.array(font_ids, dtype=np.int32), - fonts, - np.array(char_indices, dtype=np.int64), - ) - - -def get_lines_from_chars( - chars: list[Char], - superscript_height_threshold: float = 0.7, - line_distance_threshold: float = 0.1, -) -> list[Line]: - """从已提取的字符构建 pdftext lines，避免重复读取 PDFium textpage。""" - chars = _legacy_chars_to_page_chars(chars) - spans = get_spans( - chars, - superscript_height_threshold=superscript_height_threshold, - line_distance_threshold=line_distance_threshold, - ) - lines = get_lines(spans) - assign_scripts( - lines, - height_threshold=superscript_height_threshold, - line_distance_threshold=line_distance_threshold, - ) - return lines - - def _page_to_image(page: pdfium.PdfPage, scale: float, max_edge: int) -> PDFPageImage: long_edge_length = max(*page.get_size()) if (long_edge_length * scale) > max_edge: @@ -2361,3 +2252,43 @@ def _page_to_image(page: pdfium.PdfPage, scale: float, max_edge: int) -> PDFPage _try_close(bitmap) return PDFPageImage(pil_image=pil_image, scale=scale) + + +def _extract_page_text_geometry( + page: pdfium.PdfPage, + *, + include_extended_geometry: bool, +) -> PDFPageTextGeometry: + """在调用方持有的页面和锁内读取字符，使批量提取与独立接口共用实现。""" + textpage = None + try: + textpage = page.get_textpage() + raw_page_bbox: list[float] = list(page.get_bbox()) + page_bbox = _normalize_pdf_page_bbox(tuple(raw_page_bbox)) + page_rotation: int = 0 + try: + page_rotation = page.get_rotation() + except Exception: + pass + chars = get_chars(textpage, raw_page_bbox, page_rotation) + chars = _deduplicate_pdftext_chars(chars) + chars = _ensure_legacy_chars(chars) + chars = _restore_pdfium_surrogate_pairs(chars, textpage) + chars = _deduplicate_near_identical_chars(chars) + if include_extended_geometry: + loose_bboxes, tight_bboxes, origins = _extract_page_char_extended_geometry( + textpage, + chars, + page_bbox, + page_rotation, + ) + else: + loose_bboxes, tight_bboxes, origins = {}, {}, {} + finally: + _try_close(textpage) + return PDFPageTextGeometry( + chars=chars, + tight_bboxes=tight_bboxes, + origins=origins, + loose_bboxes=loose_bboxes, + ) diff --git a/mineru/model/flash/pdf/native_text.py b/mineru/model/flash/pdf/native_text.py index 1ba8ab50a..ca54efc3e 100644 --- a/mineru/model/flash/pdf/native_text.py +++ b/mineru/model/flash/pdf/native_text.py @@ -15,7 +15,7 @@ from pdftext.schema import Char from ....types import BBox -from .document import PDFDocument +from .document import PDFDocument, PDFDrawingLine from .models import ( _AxisLine, @@ -980,8 +980,14 @@ def _normalize_pdftext_angle(value: Any) -> int: def _get_pdf_drawing_lines(pdf_doc: PDFDocument, page_idx: int) -> list[_AxisLine]: """读取 PDFDocument 的公共绘图线结果，并隔离具体 PDFium 类型。""" + return _coerce_pdf_drawing_lines(pdf_doc.get_page_drawing_lines(page_idx)) + + +def _coerce_pdf_drawing_lines(drawing_lines: Sequence[PDFDrawingLine]) -> list[_AxisLine]: + """把独立接口或批量快照中的绘图线统一转换为 Flash 内部坐标类型。""" + output: list[_AxisLine] = [] - for drawing_line in pdf_doc.get_page_drawing_lines(page_idx): + for drawing_line in drawing_lines: bbox = _coerce_bbox(drawing_line.bbox) if bbox is None: continue diff --git a/mineru/model/flash/pdf/pdftext_adapter.py b/mineru/model/flash/pdf/pdftext_adapter.py new file mode 100644 index 000000000..35a6b5c8a --- /dev/null +++ b/mineru/model/flash/pdf/pdftext_adapter.py @@ -0,0 +1,175 @@ +# Copyright (c) Opendatalab. All rights reserved. +"""集中隔离 pdftext 0.6 列表与 0.7 PageChars 容器之间的适配。""" + +from __future__ import annotations + +from typing import Any, cast + +import numpy as np +from pdftext.pdf.chars import deduplicate_chars +from pdftext.pdf.pages import assign_scripts, get_lines, get_spans +from pdftext.schema import Bbox, Char, Line + +try: + from pdftext.pdf.chars import PageChars +except ImportError: + PageChars = None + + +def _is_pdftext_page_chars(chars: Any) -> bool: + """判断对象是否为 pdftext 0.7 引入的 PageChars 列式字符容器。""" + return PageChars is not None and isinstance(chars, PageChars) + + +def _deduplicate_pdftext_chars(chars: Any) -> Any: + """按当前 pdftext 返回类型调用官方去重，兼容测试或旧版本的 list 字符。""" + if _is_pdftext_page_chars(chars) or PageChars is None: + return deduplicate_chars(chars) + return chars + + +def _materialize_page_chars(chars: Any) -> list[Char]: + """将 pdftext 0.7 的 PageChars 物化为 MinerU 既有 char dict 列表。""" + boxes = chars.boxes.tolist() + rotations = chars.rotations.tolist() + font_ids = chars.font_ids.tolist() + char_indices = chars.char_indices.tolist() + + return [ + cast( + Char, + { + "bbox": Bbox([float(coord) for coord in boxes[index]]), + "char": chars.text[index], + "rotation": float(rotations[index]), + "font": chars.fonts[int(font_ids[index])], + "char_idx": int(char_indices[index]), + }, + ) + for index in range(len(chars)) + ] + + +def _ensure_legacy_chars(chars: Any) -> list[Char]: + """统一输出旧版 char dict 列表，隔离 pdftext 0.7 的返回结构变化。""" + if _is_pdftext_page_chars(chars): + return _materialize_page_chars(chars) + return cast(list[Char], chars) + + +def _char_bbox_values(bbox: object) -> tuple[float, float, float, float] | None: + """将 tuple/list 或 pdftext Bbox 对象统一转换为四元组坐标。""" + if bbox is None: + return None + if isinstance(bbox, (list, tuple)) and len(bbox) == 4: + return tuple(float(value) for value in bbox) # type: ignore[return-value] + + attrs = ("x_start", "y_start", "x_end", "y_end") + if all(hasattr(bbox, attr) for attr in attrs): + return tuple(float(getattr(bbox, attr)) for attr in attrs) # type: ignore[return-value] + return None + + +def _get_single_char_text(char: Char) -> str: + """提取单个 PDF 字符文本，异常空值用替换符保证 PageChars 长度一致。""" + text = str(char.get("char", "")) + if len(text) == 1: + return text + return text[:1] or "\ufffd" + + +def _get_char_font_id( + char: Char, + fonts: list[dict[str, Any]], + font_cache: dict[tuple[Any, Any, Any, Any], int], +) -> int: + """为旧版字符 font 生成 PageChars 需要的页内 font id。""" + font = char.get("font") or {} + font_key = ( + font.get("name"), + font.get("flags"), + font.get("size"), + font.get("weight"), + ) + font_id = font_cache.get(font_key) + if font_id is None: + font_id = len(fonts) + font_cache[font_key] = font_id + fonts.append( + { + "name": font.get("name"), + "flags": font.get("flags"), + "size": font.get("size"), + "weight": font.get("weight"), + } + ) + return font_id + + +def _get_char_index(char: Char, fallback_idx: int) -> int: + """提取旧版字符索引，缺失或为空时回退到当前列表位置。""" + char_idx = char.get("char_idx") + if char_idx is None: + char_idx = fallback_idx + return int(char_idx) + + +def _legacy_chars_to_page_chars(chars: Any) -> Any: + """将旧版 char dict 列表打包回 pdftext 0.7 get_spans 所需的 PageChars。""" + if PageChars is None or _is_pdftext_page_chars(chars): + return chars + + fonts: list[dict[str, Any]] = [] + font_cache: dict[tuple[Any, Any, Any, Any], int] = {} + text_parts: list[str] = [] + codes: list[int] = [] + rotations: list[float] = [] + boxes: list[tuple[float, float, float, float]] = [] + font_ids: list[int] = [] + char_indices: list[int] = [] + + for fallback_idx, char in enumerate(cast(list[Char], chars)): + char_text = _get_single_char_text(char) + bbox_values = _char_bbox_values(char.get("bbox")) + if bbox_values is None: + bbox_values = (0.0, 0.0, 0.0, 0.0) + text_parts.append(char_text) + codes.append(ord(char_text)) + rotations.append(float(char.get("rotation") or 0.0)) + boxes.append(bbox_values) + font_ids.append(_get_char_font_id(char, fonts, font_cache)) + char_indices.append(_get_char_index(char, fallback_idx)) + + return PageChars( + "".join(text_parts), + np.array(codes, dtype=np.uint32), + np.array(rotations, dtype=np.float64), + np.array(boxes, dtype=np.float64).reshape((len(boxes), 4)), + np.array(font_ids, dtype=np.int32), + fonts, + np.array(char_indices, dtype=np.int64), + ) + + +def get_lines_from_chars( + chars: list[Char], + superscript_height_threshold: float = 0.7, + line_distance_threshold: float = 0.1, +) -> list[Line]: + """从已提取的字符构建 pdftext lines，避免重复读取 PDFium textpage。""" + chars = _legacy_chars_to_page_chars(chars) + spans = get_spans( + chars, + superscript_height_threshold=superscript_height_threshold, + line_distance_threshold=line_distance_threshold, + ) + lines = get_lines(spans) + assign_scripts( + lines, + height_threshold=superscript_height_threshold, + line_distance_threshold=line_distance_threshold, + ) + return lines + + +__all__ = ["get_lines_from_chars"] diff --git a/mineru/model/flash/pdf/pipeline.py b/mineru/model/flash/pdf/pipeline.py index 7f1a2c5a3..b6d7264d9 100644 --- a/mineru/model/flash/pdf/pipeline.py +++ b/mineru/model/flash/pdf/pipeline.py @@ -47,7 +47,7 @@ from .native_text import ( _build_native_line_items, _extract_decorative_text_rules, - _get_pdf_drawing_lines, + _coerce_pdf_drawing_lines, _median_native_glyph_width, _sanitize_pdf_control_text, _resplit_native_visual_runs, @@ -301,28 +301,25 @@ def _analyze_native_document( ) -> list[list[dict[str, Any]]]: """逐页读取数字 PDF，并在轻量页面上完成跨页文本类型判定。""" - page_sizes = [pdf_doc.page_size(page_idx) for page_idx in range(pdf_doc.page_count)] - page_image_infos = [pdf_doc.get_page_image_infos(page_idx) for page_idx in range(pdf_doc.page_count)] - page_signature_bboxes = [pdf_doc.get_page_signature_bboxes(page_idx) for page_idx in range(pdf_doc.page_count)] - watermark_fingerprints = _detect_repeated_raster_watermark_fingerprints( - page_image_infos, - page_sizes, - ) - + page_sizes: list[tuple[float, float]] = [] + page_image_infos: list[list[PDFImageInfo]] = [] page_sources: list[_PageSource] = [] page_text_geometries = [] page_style_lines: list[list[PDFTextStyleLine]] = [] page_link_lines: list[list[PDFTextLinkLine]] = [] for page_idx in range(pdf_doc.page_count): - page_size = page_sizes[page_idx] - text_geometry = pdf_doc.get_page_chars_with_geometry(page_idx) + snapshot = pdf_doc._extract_native_page(page_idx) + page_size = snapshot.page_size + page_sizes.append(page_size) + page_image_infos.append(snapshot.image_infos) + text_geometry = snapshot.text_geometry chars = text_geometry.chars lines = _build_native_line_items( get_lines_from_chars(chars), page_size, - page_rotation=pdf_doc.page_rotation(page_idx), + page_rotation=snapshot.rotation, ) - drawing_lines = _get_pdf_drawing_lines(pdf_doc, page_idx) + drawing_lines = _coerce_pdf_drawing_lines(snapshot.drawing_lines) lines, decorative_rules = _extract_decorative_text_rules( lines, page_size, @@ -332,7 +329,7 @@ def _analyze_native_document( page_link_lines.append( detect_pdf_text_link_lines( lines, - pdf_doc.get_page_link_annotations(page_idx), + snapshot.link_annotations, ) ) source = _PageSource( @@ -340,18 +337,19 @@ def _analyze_native_document( lines=lines, chars=chars, drawing_lines=drawing_lines, - image_bboxes=_filter_repeated_raster_watermark_bboxes( - page_image_infos[page_idx], - page_size, - watermark_fingerprints, - ), - signature_bboxes=page_signature_bboxes[page_idx], - form_bboxes=pdf_doc.get_page_form_bboxes(page_idx), - path_infos=pdf_doc.get_page_path_infos(page_idx), + signature_bboxes=snapshot.signature_bboxes, + form_bboxes=snapshot.form_bboxes, + path_infos=snapshot.path_infos, ) page_sources.append(source) page_text_geometries.append(text_geometry) + watermark_fingerprints = _detect_repeated_raster_watermark_fingerprints(page_image_infos, page_sizes) + for source, image_infos in zip(page_sources, page_image_infos, strict=True): + source.image_bboxes = _filter_repeated_raster_watermark_bboxes( + image_infos, source.page_size, watermark_fingerprints, + ) + geometry_plan = build_document_geometry_plan( [source.lines for source in page_sources], page_text_geometries, diff --git a/tests/unittest/test_pdf_document.py b/tests/unittest/test_pdf_document.py index 78475af16..6acbc524a 100644 --- a/tests/unittest/test_pdf_document.py +++ b/tests/unittest/test_pdf_document.py @@ -1,5 +1,7 @@ from __future__ import annotations +from collections.abc import Callable, Iterator +from contextlib import contextmanager from io import BytesIO from typing import Any from unittest.mock import MagicMock @@ -1117,3 +1119,93 @@ def flaky_extract(*args: Any, **kwargs: Any) -> pdf_document.PDFPathInfo | None: assert call_count >= 8 assert len(path_infos) == 6 assert all(item.source_index != 0 for item in path_infos) + + +@pytest.mark.parametrize( + "builder", + [ + _build_drawing_pdf, + _build_rotated_cropped_drawing_pdf, + _build_colored_path_pdf, + _build_rotated_cropped_image_pdf, + _build_rotated_cropped_signature_pdf, + _build_rotated_cropped_link_pdf, + ], +) +def test_native_page_snapshot_matches_independent_accessors(builder: Callable[[], bytes]) -> None: + """批量提取与独立接口在旋转、裁剪、Form、签名和链接语料上完全一致。""" + + with pdf_document.PDFDocument(builder()) as document: + snapshot = document._extract_native_page(0) + assert snapshot.page_size == document.page_size(0) + assert snapshot.rotation == document.page_rotation(0) + assert snapshot.drawing_lines == document.get_page_drawing_lines(0) + assert snapshot.path_infos == document.get_page_path_infos(0) + assert snapshot.image_infos == document.get_page_image_infos(0) + assert snapshot.form_bboxes == document.get_page_form_bboxes(0) + assert snapshot.signature_bboxes == document.get_page_signature_bboxes(0) + assert snapshot.link_annotations == document.get_page_link_annotations(0) + geometry = document.get_page_chars_with_geometry(0) + assert snapshot.text_geometry.tight_bboxes == geometry.tight_bboxes + assert snapshot.text_geometry.loose_bboxes == geometry.loose_bboxes + assert snapshot.text_geometry.origins == geometry.origins + assert [(char["char"], char["char_idx"], tuple(char["bbox"])) for char in snapshot.text_geometry.chars] == [ + (char["char"], char["char_idx"], tuple(char["bbox"])) for char in geometry.chars + ] + + +def test_native_page_snapshot_opens_once_and_closes_after_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """批量读取只打开一次页面，字符提取抛错时同样关闭原生页面。""" + + opened: list[pdf_document.pdfium.PdfPage] = [] + original = pdf_document.PDFDocument._open_page + + @contextmanager + def record_page(document: pdf_document.PDFDocument, page_idx: int) -> Iterator[pdf_document.pdfium.PdfPage]: + """记录真实页面对象，使用原生命周期管理检查成功和失败后的关闭状态。""" + + with original(document, page_idx) as page: + opened.append(page) + yield page + + monkeypatch.setattr(pdf_document.PDFDocument, "_open_page", record_page) + with pdf_document.PDFDocument(_build_drawing_pdf()) as document: + document._extract_native_page(0) + assert len(opened) == 1 + assert opened[0].raw is None + + def fail_text( + page: pdf_document.pdfium.PdfPage, *, include_extended_geometry: bool + ) -> pdf_document.PDFPageTextGeometry: + """模拟字符提取失败，验证页面生命周期仍由外层上下文管理。""" + + raise RuntimeError("broken text") + + monkeypatch.setattr(pdf_document, "_extract_page_text_geometry", fail_text) + with pytest.raises(RuntimeError, match="broken text"): + document._extract_native_page(0) + assert len(opened) == 2 + assert opened[1].raw is None + + +def test_native_page_snapshot_decodes_each_path_once(monkeypatch: pytest.MonkeyPatch) -> None: + """同一 Path 解码由绘图线和路径信息共享，避免独立接口的重复工作。""" + + counts: list[object] = [] + original = pdf_document._read_raw_path_subpaths + + def record_decode(raw_object: Any) -> list[pdf_document._PathSubpath]: + """统计真实 Path 解码次数，不替换解码结果。""" + + counts.append(raw_object) + return original(raw_object) + + monkeypatch.setattr(pdf_document, "_read_raw_path_subpaths", record_decode) + with pdf_document.PDFDocument(_build_drawing_pdf()) as document: + document.get_page_drawing_lines(0) + document.get_page_path_infos(0) + separate_count = len(counts) + counts.clear() + document._extract_native_page(0) + assert separate_count > 0 + assert len(counts) * 2 == separate_count From 88bfad11770e5dbfd051663999e7fad10c2196c5 Mon Sep 17 00:00:00 2001 From: myhloli Date: Sat, 5 Sep 2026 14:42:50 +0800 Subject: [PATCH 2/2] test: count invisible paths correctly in shared PDF extraction --- tests/unittest/test_pdf_document.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unittest/test_pdf_document.py b/tests/unittest/test_pdf_document.py index 6acbc524a..5068923a9 100644 --- a/tests/unittest/test_pdf_document.py +++ b/tests/unittest/test_pdf_document.py @@ -1207,5 +1207,8 @@ def record_decode(raw_object: Any) -> list[pdf_document._PathSubpath]: separate_count = len(counts) counts.clear() document._extract_native_page(0) - assert separate_count > 0 - assert len(counts) * 2 == separate_count + shared_count = len(counts) + with document._open_page(0) as page: + expected_count = len(list(pdf_document._iter_raw_path_objects_with_depth(page))) + assert shared_count == expected_count + assert shared_count < separate_count