"""Parser for PDF files""" import hashlib import logging from io import BytesIO from pathlib import Path from typing import Dict, Iterator, List, Optional, Tuple, Union from docling_core.types.doc.base import BoundingBox, CoordOrigin from docling_core.types.doc.page import ( BitmapResource, BoundingRectangle, Coord2D, ParsedPdfDocument, PdfLine, PdfMetaData, PdfPageBoundaryType, PdfPageGeometry, PdfTableOfContents, PdfTextCell, SegmentedPdfPage, TextCell, TextDirection, ) from docling_parse.pdf_parsers import pdf_parser_v2 # type: ignore[import] from docling_parse.pdf_parsers import pdf_sanitizer # type: ignore[import] # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) class PdfDocument: def iterate_pages( self, *, keep_chars: bool = True, keep_lines: bool = True, keep_bitmaps: bool = True, create_words: bool = True, create_textlines: bool = True, enforce_same_font: bool = True, ) -> Iterator[Tuple[int, SegmentedPdfPage]]: for page_no in range(self.number_of_pages()): yield page_no + 1, self.get_page( page_no + 1, keep_chars=keep_chars, keep_lines=keep_lines, keep_bitmaps=keep_bitmaps, create_words=create_words, create_textlines=create_textlines, enforce_same_font=enforce_same_font, ) def __init__( self, parser: "pdf_parser_v2", key: str, boundary_type: PdfPageBoundaryType = PdfPageBoundaryType.CROP_BOX, ): self._parser: pdf_parser_v2 = parser self._key = key self._boundary_type = boundary_type self._pages: Dict[int, SegmentedPdfPage] = {} self._toc: Optional[PdfTableOfContents] = None self._meta: Optional[PdfMetaData] = None def is_loaded(self) -> bool: return self._parser.is_loaded(key=self._key) def unload(self) -> bool: self._pages.clear() if self.is_loaded(): return self._parser.unload_document(self._key) else: return False def unload_pages(self, page_range: tuple[int, int]): """unload page in range [page_range[0], page_range[1][""" for page_no in range(page_range[0], page_range[1]): if page_no in self._pages: self._parser.unload_document_page(key=self._key, page=page_no) del self._pages[page_no] def number_of_pages(self) -> int: if self.is_loaded(): return self._parser.number_of_pages(key=self._key) else: raise RuntimeError("This document is not loaded.") def get_meta(self) -> Optional[PdfMetaData]: if self._meta is not None: return self._meta if self.is_loaded(): xml = self._parser.get_meta_xml(key=self._key) if xml is None: return self._meta if isinstance(xml, str): self._meta = PdfMetaData(xml=xml) self._meta.initialise() return self._meta else: raise RuntimeError("This document is not loaded.") def get_table_of_contents(self) -> Optional[PdfTableOfContents]: if self.is_loaded(): toc = self._parser.get_table_of_contents(key=self._key) if toc is None: return self._toc if self._toc is not None: return self._toc self._toc = PdfTableOfContents(text="") self._toc.children = self._to_table_of_contents(toc=toc) return self._toc else: raise RuntimeError("This document is not loaded.") def _to_table_of_contents(self, toc: dict) -> List[PdfTableOfContents]: result = [] for item in toc: subtoc = PdfTableOfContents(text=item["title"]) if "children" in item: subtoc.children = self._to_table_of_contents(toc=item["children"]) result.append(subtoc) return result def get_page( self, page_no: int, *, keep_chars: bool = True, keep_lines: bool = True, keep_bitmaps: bool = True, create_words: bool = True, create_textlines: bool = True, enforce_same_font: bool = True, do_sanitization: bool = False, ) -> SegmentedPdfPage: if page_no in self._pages.keys(): return self._pages[page_no] else: if 1 <= page_no <= self.number_of_pages(): doc_dict = self._parser.parse_pdf_from_key_on_page( key=self._key, page=page_no - 1, page_boundary=self._boundary_type, do_sanitization=do_sanitization, keep_char_cells=keep_chars, keep_lines=keep_lines, keep_bitmaps=keep_bitmaps, create_word_cells=create_words, create_line_cells=create_textlines, ) for pi, page in enumerate( doc_dict["pages"] ): # only one page is expected self._pages[page_no] = self._to_segmented_page( page=page["original"], keep_chars=keep_chars, keep_lines=keep_lines, keep_bitmaps=keep_bitmaps, create_words=create_words, create_textlines=create_textlines, enforce_same_font=enforce_same_font, ) # put on cache return self._pages[page_no] raise ValueError( f"incorrect page_no: {page_no} for key={self._key} (min:1, max:{self.number_of_pages()})" ) return SegmentedPdfPage() def load_all_pages(self, create_words: bool = True, create_lines: bool = True): doc_dict = self._parser.parse_pdf_from_key( key=self._key, page_boundary=self._boundary_type, do_sanitization=False ) for pi, page in enumerate(doc_dict["pages"]): assert "original" in page, "'original' in page" # will need to be changed once we remove the original/sanitized from C++ self._pages[pi + 1] = self._to_segmented_page( page["original"], create_words=create_words, create_textlines=create_lines, ) # put on cache def _to_page_geometry(self, dimension: dict) -> PdfPageGeometry: boundary_type: PdfPageBoundaryType = PdfPageBoundaryType( dimension["page_boundary"] ) art_bbox = BoundingBox( l=dimension["rectangles"]["art-bbox"][0], b=dimension["rectangles"]["art-bbox"][1], r=dimension["rectangles"]["art-bbox"][2], t=dimension["rectangles"]["art-bbox"][3], coord_origin=CoordOrigin.BOTTOMLEFT, ) media_bbox = BoundingBox( l=dimension["rectangles"]["media-bbox"][0], b=dimension["rectangles"]["media-bbox"][1], r=dimension["rectangles"]["media-bbox"][2], t=dimension["rectangles"]["media-bbox"][3], coord_origin=CoordOrigin.BOTTOMLEFT, ) bleed_bbox = BoundingBox( l=dimension["rectangles"]["bleed-bbox"][0], b=dimension["rectangles"]["bleed-bbox"][1], r=dimension["rectangles"]["bleed-bbox"][2], t=dimension["rectangles"]["bleed-bbox"][3], coord_origin=CoordOrigin.BOTTOMLEFT, ) trim_bbox = BoundingBox( l=dimension["rectangles"]["trim-bbox"][0], b=dimension["rectangles"]["trim-bbox"][1], r=dimension["rectangles"]["trim-bbox"][2], t=dimension["rectangles"]["trim-bbox"][3], coord_origin=CoordOrigin.BOTTOMLEFT, ) crop_bbox = BoundingBox( l=dimension["rectangles"]["crop-bbox"][0], b=dimension["rectangles"]["crop-bbox"][1], r=dimension["rectangles"]["crop-bbox"][2], t=dimension["rectangles"]["crop-bbox"][3], coord_origin=CoordOrigin.BOTTOMLEFT, ) # Fixme: The boundary type to which this rect refers should accept a user argument # TODO: Why is this a BoundingRectangle not a BoundingBox? rect = BoundingRectangle( r_x0=crop_bbox.l, r_y0=crop_bbox.b, r_x1=crop_bbox.r, r_y1=crop_bbox.b, r_x2=crop_bbox.r, r_y2=crop_bbox.t, r_x3=crop_bbox.l, r_y3=crop_bbox.t, coord_origin=CoordOrigin.BOTTOMLEFT, ) return PdfPageGeometry( angle=dimension["angle"], boundary_type=boundary_type, rect=rect, art_bbox=art_bbox, media_bbox=media_bbox, trim_bbox=trim_bbox, crop_bbox=crop_bbox, bleed_bbox=bleed_bbox, ) def _to_cells(self, cells: dict) -> List[Union[PdfTextCell, TextCell]]: assert "data" in cells, '"data" in cells' assert "header" in cells, '"header" in cells' data = cells["data"] header = cells["header"] # Pre-compute header indices as local variables r_x0_idx = header.index("r_x0") r_y0_idx = header.index("r_y0") r_x1_idx = header.index("r_x1") r_y1_idx = header.index("r_y1") r_x2_idx = header.index("r_x2") r_y2_idx = header.index("r_y2") r_x3_idx = header.index("r_x3") r_y3_idx = header.index("r_y3") text_idx = header.index("text") font_key_idx = header.index("font-key") font_name_idx = header.index("font-name") widget_idx = header.index("widget") left_to_right_idx = header.index("left_to_right") rendering_mode_idx = header.index("rendering-mode") # Pre-allocate list with exact size data_len = len(data) result: List[Union[PdfTextCell, TextCell]] = [None] * data_len # type: ignore for ind, row in enumerate(data): rect = BoundingRectangle( r_x0=row[r_x0_idx], r_y0=row[r_y0_idx], r_x1=row[r_x1_idx], r_y1=row[r_y1_idx], r_x2=row[r_x2_idx], r_y2=row[r_y2_idx], r_x3=row[r_x3_idx], r_y3=row[r_y3_idx], ) result[ind] = PdfTextCell( rect=rect, text=row[text_idx], orig=row[text_idx], font_key=row[font_key_idx], font_name=row[font_name_idx], widget=row[widget_idx], text_direction=( TextDirection.LEFT_TO_RIGHT if row[left_to_right_idx] else TextDirection.RIGHT_TO_LEFT ), index=ind, rendering_mode=row[rendering_mode_idx], ) return result def _to_bitmap_resources(self, images: dict) -> List[BitmapResource]: assert "data" in images, '"data" in images' assert "header" in images, '"header" in images' data = images["data"] header = images["header"] result: List[BitmapResource] = [] for ind, row in enumerate(data): rect = BoundingRectangle( r_x0=row[header.index(f"x0")], r_y0=row[header.index(f"y0")], r_x1=row[header.index(f"x1")], r_y1=row[header.index(f"y0")], r_x2=row[header.index(f"x1")], r_y2=row[header.index(f"y1")], r_x3=row[header.index(f"x0")], r_y3=row[header.index(f"y1")], ) image = BitmapResource(index=ind, rect=rect, uri=None) result.append(image) return result def _to_lines(self, data: dict) -> List[PdfLine]: result: List[PdfLine] = [] for ind, item in enumerate(data): for l in range(0, len(item["i"]), 2): i0: int = item["i"][l + 0] i1: int = item["i"][l + 1] points: List[Coord2D] = [] for k in range(i0, i1): points.append(Coord2D(item["x"][k], item["y"][k])) line = PdfLine( index=ind, parent_id=l, points=points, ) result.append(line) return result def _to_segmented_page( self, page: dict, *, keep_chars: bool = True, keep_lines: bool = True, keep_bitmaps: bool = True, create_words: bool, create_textlines: bool, enforce_same_font: bool = True, ) -> SegmentedPdfPage: # FIXME: this might be inefficient ... """ char_cells = self._to_cells(page["cells"]) segmented_page = SegmentedPdfPage( dimension=self._to_page_geometry(page["dimension"]), char_cells=char_cells, word_cells=[], textline_cells=[], has_chars=len(char_cells) > 0, bitmap_resources=self._to_bitmap_resources(page["images"]), lines=self._to_lines(page["lines"]), ) if create_words: self._create_word_cells(segmented_page, enforce_same_font=enforce_same_font) if create_textlines: self._create_textline_cells( segmented_page, enforce_same_font=enforce_same_font ) """ char_cells = [] if keep_chars: assert "cells" in page char_cells = self._to_cells(page["cells"]) lines = [] if keep_lines: assert "lines" in page lines = self._to_lines(page["lines"]) bitmap_resources = [] if keep_bitmaps: assert "images" in page bitmap_resources = self._to_bitmap_resources(page["images"]) segmented_page = SegmentedPdfPage( dimension=self._to_page_geometry(page["dimension"]), char_cells=char_cells, word_cells=[], textline_cells=[], has_chars=len(char_cells) > 0, bitmap_resources=bitmap_resources, # self._to_bitmap_resources(page["images"]), lines=lines, # self._to_lines(page["lines"]), ) if create_words and ("word_cells" in page): segmented_page.word_cells = self._to_cells(page["word_cells"]) segmented_page.has_words = len(segmented_page.word_cells) > 0 elif keep_chars: logging.warning( "`words` will be created for segmented_page in an inefficient way!" ) self._create_word_cells(segmented_page, enforce_same_font=enforce_same_font) # else: # logging.warning("No `words` will be created for segmented_page") if create_textlines and ("line_cells" in page): segmented_page.textline_cells = self._to_cells(page["line_cells"]) segmented_page.has_lines = len(segmented_page.textline_cells) > 0 elif keep_chars: logging.warning( "`text_lines` will be created for segmented_page in an inefficient way!" ) self._create_textline_cells( segmented_page, enforce_same_font=enforce_same_font ) # else: # logging.warning("No `text_lines` will be created for segmented_page") return segmented_page def _create_word_cells( self, segmented_page: SegmentedPdfPage, *, horizontal_cell_tolerance: float = 1.0, space_width_factor_for_merge: float = 0.33, enforce_same_font: bool = True, _loglevel: str = "fatal", ): if len(segmented_page.word_cells) > 0: return sanitizer = pdf_sanitizer(level=_loglevel) char_data = [] for item in segmented_page.char_cells: item_dict = item.model_dump(mode="json", by_alias=True, exclude_none=True) item_dict["left_to_right"] = ( item.text_direction == TextDirection.LEFT_TO_RIGHT ) char_data.append(item_dict) sanitizer.set_char_cells(data=char_data) # data = sanitizer.create_word_cells(space_width_factor_for_merge=0.33) data = sanitizer.create_word_cells( horizontal_cell_tolerance=horizontal_cell_tolerance, space_width_factor_for_merge=space_width_factor_for_merge, enforce_same_font=enforce_same_font, ) segmented_page.word_cells = [] for item in data: cell = PdfTextCell.model_validate(item) segmented_page.word_cells.append(cell) segmented_page.has_words = len(segmented_page.word_cells) > 0 def _create_textline_cells( self, segmented_page: SegmentedPdfPage, *, horizontal_cell_tolerance: float = 1.0, space_width_factor_for_merge: float = 1.0, space_width_factor_for_merge_with_space: float = 0.33, enforce_same_font: bool = True, _loglevel: str = "fatal", ): if len(segmented_page.textline_cells) > 0: return sanitizer = pdf_sanitizer(level=_loglevel) char_data = [] for item in segmented_page.char_cells: item_dict = item.model_dump(mode="json", by_alias=True, exclude_none=True) # TODO changing representation for the C++ parser, need to update on C++ code. item_dict["left_to_right"] = ( item.text_direction == TextDirection.LEFT_TO_RIGHT ) item_dict["id"] = item.index char_data.append(item_dict) sanitizer.set_char_cells(data=char_data) # data = sanitizer.create_line_cells() data = sanitizer.create_line_cells( horizontal_cell_tolerance=horizontal_cell_tolerance, space_width_factor_for_merge=space_width_factor_for_merge, space_width_factor_for_merge_with_space=space_width_factor_for_merge_with_space, enforce_same_font=enforce_same_font, ) segmented_page.textline_cells = [] for item in data: cell = PdfTextCell.model_validate(item) segmented_page.textline_cells.append(cell) segmented_page.has_lines = len(segmented_page.textline_cells) > 0 def _to_parsed_document( self, doc_dict: dict, page_no: int = 1, create_words: bool = False, create_lines: bool = True, ) -> ParsedPdfDocument: parsed_doc = ParsedPdfDocument() for pi, page in enumerate(doc_dict["pages"]): parsed_doc.pages[page_no + pi] = self._to_segmented_page( page["original"], create_words=create_words, create_textlines=create_lines, ) return parsed_doc class DoclingPdfParser: def __init__(self, loglevel: str = "fatal"): """ Set the log level using a string label. Parameters: level (str): Logging level as a string. One of ['fatal', 'error', 'warning', 'info'] """ self.parser = pdf_parser_v2(level=loglevel) def set_loglevel(self, loglevel: str): """Set the log level using a string label. Parameters: level (str): Logging level as a string. One of ['fatal', 'error', 'warning', 'info'] )") """ self.parser.set_loglevel_with_label(level=loglevel) def list_loaded_keys(self) -> List[str]: """List the keys of the loaded documents. Returns: List[str]: A list of keys for the currently loaded documents. """ return self.parser.list_loaded_keys() def load( self, path_or_stream: Union[str, Path, BytesIO], lazy: bool = True, boundary_type: PdfPageBoundaryType = PdfPageBoundaryType.CROP_BOX, password: Optional[str] = None, ) -> PdfDocument: if isinstance(path_or_stream, str): path_or_stream = Path(path_or_stream) if isinstance(path_or_stream, Path): key = f"key={str(path_or_stream)}" # use filepath as internal handle success = self._load_document( key=key, filename=str(path_or_stream), password=password ) elif isinstance(path_or_stream, BytesIO): hasher = hashlib.sha256(usedforsecurity=False) while chunk := path_or_stream.read(8192): hasher.update(chunk) path_or_stream.seek(0) hash = hasher.hexdigest() key = f"key={hash}" # use md5 hash as internal handle success = self._load_document_from_bytesio(key=key, data=path_or_stream) if success: result_doc = PdfDocument( parser=self.parser, key=key, boundary_type=boundary_type ) if not lazy: # eagerly parse the pages at init time if desired result_doc.load_all_pages() return result_doc else: raise RuntimeError(f"Failed to load document with key {key}") def _load_document( self, key: str, filename: str, password: Optional[str] = None ) -> bool: """Load a document by key and filename. Parameters: key (str): The unique key to identify the document. filename (str): The path to the document file to load. password (str, optional): Optional password for password-protected files Returns: bool: True if the document was successfully loaded, False otherwise.)") """ return self.parser.load_document( key=key, filename=filename.encode("utf8"), password=password ) def _load_document_from_bytesio(self, key: str, data: BytesIO) -> bool: """Load a document by key from a BytesIO-like object. Parameters: key (str): The unique key to identify the document. bytes_io (Any): A BytesIO-like object containing the document data. Returns: bool: True if the document was successfully loaded, False otherwise.)") """ return self.parser.load_document_from_bytesio(key=key, bytes_io=data)