diff --git a/README.md b/README.md index fb025dda..73c45bd3 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,38 @@ Path("output.md").write_bytes(md.encode()) | **Plain text** | `to_text(path)` | Search indexing, simple NLP tasks | | **LlamaIndex docs** | `LlamaMarkdownReader().load_data(path)` | Direct LlamaIndex integration | +### HTML table output + +By default, tables are rendered as GitHub-compatible Markdown. To emit +reconstructed HTML `` elements instead, use `table_output="html"`: + +```python +import pymupdf4llm + +md = pymupdf4llm.to_markdown("document.pdf", table_output="html") +``` + +In layout mode this keeps the layout pipeline for reading order, body text, and +OCR handling, and swaps only table rendering to the HTML table engine. The +standalone table API is also available: + +```python +from pymupdf4llm.helpers.table_html import to_html + +html = to_html("document.pdf", page_index=0) +``` + +### Layout edge threshold + +Layout mode accepts `edge_threshold`, passed through to `page.get_layout()` +to adjust the layout model's grouping confidence: + +```python +md = pymupdf4llm.to_markdown("document.pdf", edge_threshold=0.75) +``` + +The default remains the library default unless this option is provided. + ### Extraction capabilities | Feature | Description | @@ -532,4 +564,3 @@ If you find this useful, please consider giving it a star — it helps others di [![Star on GitHub](https://img.shields.io/github/stars/pymupdf/pymupdf4llm.svg?style=for-the-badge&label=Star&logo=github)](https://github.com/pymupdf/pymupdf4llm/) - diff --git a/src/__init__.py b/src/__init__.py index d2f7da0b..59ef1593 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -80,6 +80,8 @@ def _layout_to_markdown( show_progress=False, use_ocr=True, write_images=False, + render_html_tables=None, + edge_threshold=None, # unsupported options for pymupdf layout: **kwargs, ): @@ -101,6 +103,8 @@ def _layout_to_markdown( force_ocr=force_ocr, ocr_language=ocr_language, ocr_function=ocr_function, + render_html_tables=render_html_tables, + edge_threshold=edge_threshold, ) return parsed_doc.to_markdown( header=header, @@ -129,6 +133,8 @@ def _layout_to_json( force_ocr=False, ocr_language="eng", ocr_function=None, + render_html_tables=None, + edge_threshold=None, # unsupported options for pymupdf layout: **kwargs, ): @@ -146,6 +152,8 @@ def _layout_to_json( force_ocr=force_ocr, ocr_language=ocr_language, ocr_function=ocr_function, + render_html_tables=render_html_tables, + edge_threshold=edge_threshold, ) return parsed_doc.to_json() @@ -168,6 +176,7 @@ def _layout_to_text( table_max_width=100, table_min_col_width=10, page_chunks=False, + edge_threshold=None, # unsupported options for pymupdf layout: **kwargs, ): @@ -183,6 +192,7 @@ def _layout_to_text( force_ocr=force_ocr, ocr_language=ocr_language, ocr_function=ocr_function, + edge_threshold=edge_threshold, ) return parsed_doc.to_text( header=header, @@ -197,6 +207,35 @@ def _layout_to_text( def to_markdown(*args, **kwargs): + # `render_html_tables` is an internal flag this wrapper injects for + # table_output="html"; it is not a public kwarg. Drop any user-supplied value + # so it cannot silently enable/disable HTML tables via **kwargs. + kwargs.pop("render_html_tables", None) + if kwargs.get("table_output") == "html": + # Render tables as HTML via table_html. + kwargs = dict(kwargs) + kwargs.pop("table_output", None) + if _use_layout: + # Preferred path: render HTML tables on the layout path, reusing the + # GNN layout. Keeps the layout path's text, reading order, and OCR -- + # only the table rendering is swapped. + return _layout_to_markdown(*args, render_html_tables=True, **kwargs) + # No layout engine available: fall back to the rag path, which has its own + # table_output="html" wiring. OCR is not available on this path. + legacy_kwargs = dict(kwargs) + for name in ( + "footer", + "header", + "ocr_dpi", + "ocr_function", + "ocr_language", + "use_ocr", + "force_ocr", + ): + legacy_kwargs.pop(name, None) + return pymupdf4llm.helpers.pymupdf_rag.to_markdown( + *args, table_output="html", **legacy_kwargs + ) if _use_layout: return _layout_to_markdown(*args, **kwargs) else: @@ -204,6 +243,12 @@ def to_markdown(*args, **kwargs): def to_json(*args, **kwargs): + # See to_markdown: `render_html_tables` is internal, not a public kwarg. + kwargs.pop("render_html_tables", None) + if kwargs.get("table_output") == "html": + kwargs = dict(kwargs) + kwargs.pop("table_output", None) + kwargs["render_html_tables"] = True if _use_layout: return _layout_to_json(*args, **kwargs) else: diff --git a/src/helpers/document_layout.py b/src/helpers/document_layout.py index 8271999e..6c63d212 100644 --- a/src/helpers/document_layout.py +++ b/src/helpers/document_layout.py @@ -1,8 +1,8 @@ import base64 import io import json -import os import math +import threading from dataclasses import dataclass from collections import defaultdict from pathlib import Path @@ -24,6 +24,7 @@ from dataclasses import dataclass pymupdf.TOOLS.unset_quad_corrections(True) +_LAYOUT_LOCK = threading.RLock() INFO_MESSAGES = io.StringIO() GRAPHICS_TEXT = "\n![](%s)\n" @@ -40,6 +41,12 @@ BULLETS = tuple(utils.BULLETS) +def get_layout_locked(page: pymupdf.Page, **kwargs): + """Serialize PyMuPDF layout inference, which uses process-global state.""" + with _LAYOUT_LOCK: + return page.get_layout(**kwargs) + + def get_table_details(tab_dict, table_blocks): """Create a TableDetails object. @@ -712,6 +719,162 @@ def fallback_text_to_md(textlines, ignore_code: bool = False, clip=None): return output + "\n" +def _rect_area(rect) -> float: + return max(0.0, float(rect.x1 - rect.x0)) * max(0.0, float(rect.y1 - rect.y0)) + + +def _html_table_meta(table_item) -> Dict: + rect = pymupdf.Rect(table_item[0]) + return { + "bbox": [float(rect.x0), float(rect.y0), float(rect.x1), float(rect.y1)], + "html": table_item[1], + "rows": int(table_item[2]) if len(table_item) > 2 and table_item[2] is not None else None, + "cols": int(table_item[3]) if len(table_item) > 3 and table_item[3] is not None else None, + "cells": table_item[4] if len(table_item) > 4 else None, + "extract": table_item[5] if len(table_item) > 5 else None, + } + + +def _assign_html_tables_to_boxes(layout_boxes, html_tables, threshold: float = 0.5): + """Assign table_html output to layout table boxes, adding unmatched boxes.""" + by_box: Dict[tuple, List[Dict]] = {} + if not html_tables: + return layout_boxes, by_box + + table_boxes = [ + (tuple(pymupdf.IRect(box[:4])), pymupdf.Rect(box[:4])) + for box in layout_boxes + if len(box) >= 5 and box[4] == "table" + ] + augmented_boxes = list(layout_boxes) + + for table_item in html_tables: + meta = _html_table_meta(table_item) + table_rect = pymupdf.Rect(meta["bbox"]) + table_area = _rect_area(table_rect) + best_key = None + best_score = 0.0 + if table_area > 0: + for key, box_rect in table_boxes: + inter = table_rect & box_rect + if inter.is_empty: + continue + score = _rect_area(inter) / table_area + if score > best_score: + best_key = key + best_score = score + if best_key is None or best_score < threshold: + synthetic = (table_rect.x0, table_rect.y0, table_rect.x1, table_rect.y1, "table") + augmented_boxes.append(synthetic) + best_key = tuple(pymupdf.IRect(table_rect)) + table_boxes.append((best_key, table_rect)) + by_box.setdefault(best_key, []).append(meta) + + for items in by_box.values(): + items.sort(key=lambda item: (item["bbox"][1], item["bbox"][0])) + return augmented_boxes, by_box + + +def _line_claimed_by_table(line_bbox, table_rects) -> bool: + line_rect = pymupdf.Rect(line_bbox) + line_area = _rect_area(line_rect) + if line_area <= 0: + return False + center = pymupdf.Point( + (line_rect.x0 + line_rect.x1) * 0.5, + (line_rect.y0 + line_rect.y1) * 0.5, + ) + for table_rect in table_rects: + if center in table_rect: + return True + inter = line_rect & table_rect + if not inter.is_empty and _rect_area(inter) / line_area >= 0.5: + return True + return False + + +def _union_line_rects(lines): + rect = pymupdf.Rect(lines[0]["bbox"]) + for line in lines[1:]: + rect |= pymupdf.Rect(line["bbox"]) + return rect + + +def _split_text_box_around_tables(box, fulltext, table_rects): + box_rect = pymupdf.Rect(box[:4]) + if not any(box_rect.intersects(table_rect) for table_rect in table_rects): + return [box], {} + + try: + lines = [ + {"bbox": l[0], "spans": l[1]} + for l in get_raw_lines( + textpage=None, + blocks=fulltext, + clip=box_rect, + ignore_invisible=False, + ) + ] + except Exception: + return [box], {} + if not lines: + return [box], {} + + split_boxes = [] + textlines_by_box = {} + claimed = [_line_claimed_by_table(line["bbox"], table_rects) for line in lines] + start = None + for index, is_claimed in enumerate(claimed + [True]): + if not is_claimed and start is None: + start = index + elif is_claimed and start is not None: + end = index - 1 + y0 = box_rect.y0 if start == 0 else pymupdf.Rect(lines[start - 1]["bbox"]).y1 + y1 = ( + box_rect.y1 + if end == len(lines) - 1 + else pymupdf.Rect(lines[end + 1]["bbox"]).y0 + ) + if y1 <= y0: + rect = _union_line_rects(lines[start : end + 1]) + y0, y1 = rect.y0, rect.y1 + split_box = (box_rect.x0, y0, box_rect.x1, y1, box[4]) + split_boxes.append(split_box) + textlines_by_box[tuple(pymupdf.IRect(split_box[:4]))] = lines[start : end + 1] + start = None + return split_boxes, textlines_by_box + + +def normalize_layout_boxes(layout_boxes, html_tables, fulltext): + """Build opt-in HTML-table layout boxes before reading-order sorting.""" + normalized_boxes, html_tables_by_box = _assign_html_tables_to_boxes( + layout_boxes, + html_tables, + ) + if not html_tables_by_box: + return normalized_boxes, html_tables_by_box, {} + + table_rects = [ + pymupdf.Rect(item["bbox"]) + for items in html_tables_by_box.values() + for item in items + ] + output_boxes = [] + textlines_by_box = {} + for box in normalized_boxes: + if len(box) < 5 or box[4] in ("table", "picture", "formula"): + output_boxes.append(box) + continue + split_boxes, split_textlines = _split_text_box_around_tables( + box, + fulltext, + table_rects, + ) + output_boxes.extend(split_boxes) + textlines_by_box.update(split_textlines) + return output_boxes, html_tables_by_box, textlines_by_box + + @dataclass class TableDetails: bbox: tuple = None @@ -833,6 +996,10 @@ def to_markdown( string_lengths.append(len(md_string)) continue if btype == "table": + if box.table.get("html"): + md_string += box.table["html"] + "\n\n" + string_lengths.append(len(md_string)) + continue table_text = box.table["markdown"] if page.full_ocred: # remove code style if page was OCR'd @@ -1091,6 +1258,8 @@ def parse_document( force_ocr=False, ocr_language="eng", ocr_function=None, + render_html_tables=None, + edge_threshold=None, ) -> ParsedDocument: if isinstance(doc, pymupdf.Document): mydoc = doc @@ -1207,17 +1376,42 @@ def parse_document( textpage = page.get_textpage(flags=FLAGS, clip=pymupdf.INFINITE_RECT()) blocks = textpage.extractDICT()["blocks"] - # Execute the Layout module AFTER any OCR - page.get_layout(return_raw=True) - - # Determine if any tables are present. If False, we skip any table-related efforts. - tables_exist = any( - b for b in page.layout_information if b["class_name"] == "table" - ) + # Execute the Layout module AFTER any OCR. + layout_kwargs = {"return_raw": True} + if edge_threshold is not None: + layout_kwargs["edge_threshold"] = edge_threshold + get_layout_locked(page, **layout_kwargs) + + # Optionally render tables as HTML, reusing this raw GNN layout + # (get_layout is guarded to reuse it, so no second GNN pass). + # Save/restore layout_information so table_html's internal normalization + # does not disturb the layout path below. + page_html_tables_list = None + _render_html_tables = bool(render_html_tables) + if _render_html_tables: + _saved_raw_layout = page.layout_information + try: + from pymupdf4llm.helpers.table_html import page_html_tables + page_html_tables_list = list(page_html_tables(page)) + except Exception as exc: + # HTML table rendering failed on this page -> fall back to the + # layout table path, but surface the reason (flushed to + # pymupdf.message at the end of parse_document) rather than + # swallowing it silently. + print( + f"Warning: HTML table engine failed on page {page.number}: " + f"{type(exc).__name__}: {exc}", + file=INFO_MESSAGES, + ) + page_html_tables_list = None + finally: + page.layout_information = _saved_raw_layout # Dictionary with details for all tables. Key is the bounding box # tuple, value is the original Layout info per table. table_infos = {} + html_tables_by_box = {} + textlines_by_box = {} new_layout_info = [] # will contain Layout boxes in non-"raw" format for b in page.layout_information: @@ -1230,7 +1424,20 @@ def parse_document( key = tuple(pymupdf.IRect(b["group_bbox"])) table_infos[key] = b + if _render_html_tables and page_html_tables_list: + new_layout_info, html_tables_by_box, textlines_by_box = normalize_layout_boxes( + new_layout_info, + page_html_tables_list, + fulltext=[b for b in blocks if b["type"] == 0], + ) + page.layout_information = new_layout_info + # Determine if any tables are present after HTML-table normalization. + # Synthetic find_tables proposals are first-class table boxes on the + # opt-in path, so this must be based on the normalized boxes. + tables_exist = any( + len(b) >= 5 and b[4] == "table" for b in page.layout_information + ) if not OCR_SPANS: # some cleaning if no old OCR spans utils.clean_pictures(page, blocks) utils.add_image_orphans(page, blocks) @@ -1296,36 +1503,74 @@ def parse_document( elif layoutbox.boxclass == "table": search_key = (layoutbox.x0, layoutbox.y0, layoutbox.x1, layoutbox.y1) - - # Because of intermediate processing, the bbox might not match - # the original exactly. So we need to take the best fit. - key = max(table_infos.keys(), key=lambda k: utils.iou(k, search_key)) - - tab_dict = table_infos.get(key) - tab_details = get_table_details(tab_dict, table_blocks) - - layoutbox.table = { - "bbox": list(tab_details.bbox), - "row_count": tab_details.row_count, - "col_count": tab_details.col_count, - "cells": tab_details.cells, - "extract": tab_details.extract, - "markdown": tab_details.markdown, - } + html_tables = html_tables_by_box.get(tuple(pymupdf.IRect(clip)), []) + + if html_tables: + # Opt-in HTML mode: this box's table(s) were detected and + # rendered by table_html. Its reconstructed grid -- not the + # layout GNN grid or a best-fit rematch -- is the source of + # truth, so row_count/col_count/cells/extract describe the + # SAME grid the html shows: `cells` the post-span cell bbox + # matrix and `extract` the parallel plain-text matrix (both + # None for span-covered slots). `markdown` stays None -- + # `html` is authoritative in this mode (to_markdown emits + # box.table["html"] directly). + single = html_tables[0] if len(html_tables) == 1 else None + layoutbox.table = { + "bbox": [layoutbox.x0, layoutbox.y0, layoutbox.x1, layoutbox.y1], + "row_count": single.get("rows") if single else None, + "col_count": single.get("cols") if single else None, + "cells": single.get("cells") if single else None, + "extract": single.get("extract") if single else None, + "markdown": None, + "html_tables": html_tables, + "html": "\n\n".join(item["html"] for item in html_tables), + } + else: + # Non-HTML path (to_text, or a layout table box table_html did + # not render): keep the layout-grid extraction that feeds + # markdown/text. Because of intermediate processing the bbox + # might not match the original exactly, so take the best fit. + tab_details = None + if table_infos and table_blocks is not None: + key = max(table_infos.keys(), key=lambda k: utils.iou(k, search_key)) + tab_dict = table_infos.get(key) + tab_details = get_table_details(tab_dict, table_blocks) + + if tab_details is not None: + layoutbox.table = { + "bbox": list(tab_details.bbox), + "row_count": tab_details.row_count, + "col_count": tab_details.col_count, + "cells": tab_details.cells, + "extract": tab_details.extract, + "markdown": tab_details.markdown, + } + else: + layoutbox.table = { + "bbox": [layoutbox.x0, layoutbox.y0, layoutbox.x1, layoutbox.y1], + "row_count": None, + "col_count": None, + "cells": None, + "extract": None, + "markdown": "", + } else: # Handle text-like box classes: # Extract text line information within the box. # Each line is represented as its bbox and a list of spans. - layoutbox.textlines = [ - {"bbox": l[0], "spans": l[1]} - for l in get_raw_lines( - textpage=None, - blocks=pagelayout.fulltext, - clip=clip, - ignore_invisible=False, - ) - ] + layoutbox.textlines = textlines_by_box.get(tuple(pymupdf.IRect(clip))) + if layoutbox.textlines is None: + layoutbox.textlines = [ + {"bbox": l[0], "spans": l[1]} + for l in get_raw_lines( + textpage=None, + blocks=pagelayout.fulltext, + clip=clip, + ignore_invisible=False, + ) + ] # For each title/section_header compute and store the maximum # font size, to be used as a signal for header "#" prefix if layoutbox.boxclass in ("title", "section-header"): diff --git a/src/helpers/pymupdf_rag.py b/src/helpers/pymupdf_rag.py index b60d3163..18860322 100644 --- a/src/helpers/pymupdf_rag.py +++ b/src/helpers/pymupdf_rag.py @@ -321,6 +321,35 @@ def to_text(*args, **kwargs): ) +# -------------------------------------------------------------------------- +# HTML table output (table_output="html") +# +# When requested, tables are detected and rendered as HTML
by +# pymupdf4llm.helpers.table_html. In this mode table_html -- not a direct +# page.find_tables() call here -- is the source of the page's tables, so it +# drives emission, reading order, and body-text exclusion. parms.tab_rects +# index i maps directly to the i-th reconstructed table in parms.html_tables. +# -------------------------------------------------------------------------- +def _reconstruct_html_tables(page): + """Return the ``(bbox, html, rows, cols, cells, extract)`` payload for each + table on ``page``. Imported lazily so table_html loads only when + table_output="html" is requested.""" + from pymupdf4llm.helpers.table_html import page_html_tables + + return page_html_tables(page) + + +def _table_string(parms, i, table_output): + """Render table ``i`` as markdown (default) or reconstructed HTML. + + In HTML mode ``i`` indexes ``parms.html_tables`` directly, so no bbox + matching is needed. + """ + if table_output == "html": + return parms.html_tables[i][1] + return parms.tabs[i].to_markdown(clean=False) + + def to_markdown( doc, *, @@ -343,6 +372,7 @@ def to_markdown( page_width=612, page_height=None, table_strategy="lines_strict", + table_output="markdown", graphics_limit=None, fontsize_limit=3, ignore_code=False, @@ -370,6 +400,8 @@ def to_markdown( page_width: (float) assumption if page layout is variable. page_height: (float) assumption if page layout is variable. table_strategy: choose table detection strategy + table_output: ("markdown" or "html") render tables as markdown (default) + or as reconstructed HTML
via pymupdf4llm.helpers.table_html. graphics_limit: (int) if vector graphics count exceeds this, ignore all. ignore_code: (bool) suppress code-like formatting (mono-space fonts) extract_words: (bool, False) include "words"-like output in page chunks @@ -395,6 +427,9 @@ def to_markdown( if EXTRACT_WORDS is True: page_chunks = True ignore_code = True + if table_output not in ("markdown", "html"): + raise ValueError("'table_output' must be 'markdown' or 'html'.") + TABLE_OUTPUT = table_output IMG_PATH = image_path if IMG_PATH and write_images is True and not os.path.exists(IMG_PATH): os.makedirs(IMG_PATH, exist_ok=True) @@ -577,8 +612,10 @@ def write_text( ) ] for i, _ in tab_candidates: - out_string += "\n" + parms.tabs[i].to_markdown(clean=False) + "\n" - if EXTRACT_WORDS: + out_string += "\n" + _table_string(parms, i, TABLE_OUTPUT) + "\n" + # HTML mode has no find_tables Table objects to read cell rects + # from, so table cell rects are not wired for "words" output yet. + if EXTRACT_WORDS and TABLE_OUTPUT != "html": # for "words" extraction, add table cells as line rects cells = sorted( set( @@ -818,8 +855,8 @@ def output_tables(parms, text_rect): ): if i in parms.written_tables: continue - this_md += parms.tabs[i].to_markdown(clean=False) + "\n" - if EXTRACT_WORDS: + this_md += _table_string(parms, i, TABLE_OUTPUT) + "\n" + if EXTRACT_WORDS and TABLE_OUTPUT != "html": # for "words" extraction, add table cells as line rects cells = sorted( set( @@ -839,8 +876,8 @@ def output_tables(parms, text_rect): for i, trect in parms.tab_rects.items(): if i in parms.written_tables: continue - this_md += parms.tabs[i].to_markdown(clean=False) + "\n" - if EXTRACT_WORDS: + this_md += _table_string(parms, i, TABLE_OUTPUT) + "\n" + if EXTRACT_WORDS and TABLE_OUTPUT != "html": # for "words" extraction, add table cells as line rects cells = sorted( set( @@ -1093,32 +1130,47 @@ def get_page_output( # Locate all tables on page parms.written_tables = [] # stores already written tables - omitted_table_rects = [] parms.tabs = [] + parms.html_tables = [] + tab_rects = {} + if IGNORE_GRAPHICS or not table_strategy: # do not try to extract tables pass + elif TABLE_OUTPUT == "html": + # table_html detects (via the layout stage, with find_tables repair) + # and reconstructs each table as HTML, and drives table emission, + # reading order, and body-text exclusion below. It calls + # page.find_tables() internally (once); to_markdown does not call it + # directly here. Borderless tables the layout stage finds -- which + # find_tables alone cannot see -- are emitted too. + parms.html_tables = _reconstruct_html_tables(page) + for i, (rect, _html, rows, cols, _cells, _extract) in enumerate(parms.html_tables): + tab_rects[i] = pymupdf.Rect(rect) + parms.tables.append( + {"bbox": tuple(tab_rects[i]), "rows": rows, "columns": cols} + ) else: tabs = page.find_tables(clip=parms.clip, strategy=table_strategy) for t in tabs.tables: # remove tables with too few rows or columns if t.row_count < 2 or t.col_count < 2: - omitted_table_rects.append(pymupdf.Rect(t.bbox)) continue parms.tabs.append(t) parms.tabs.sort(key=lambda t: (t.bbox[0], t.bbox[1])) - # Make a list of table boundary boxes. - # Must include the header bbox (which may exist outside tab.bbox) - tab_rects = {} - for i, t in enumerate(parms.tabs): - tab_rects[i] = pymupdf.Rect(t.bbox) | pymupdf.Rect(t.header.bbox) - tab_dict = { - "bbox": tuple(tab_rects[i]), - "rows": t.row_count, - "columns": t.col_count, - } - parms.tables.append(tab_dict) + # Make a list of table boundary boxes. + # Must include the header bbox (which may exist outside tab.bbox) + for i, t in enumerate(parms.tabs): + tab_rects[i] = pymupdf.Rect(t.bbox) | pymupdf.Rect(t.header.bbox) + parms.tables.append( + { + "bbox": tuple(tab_rects[i]), + "rows": t.row_count, + "columns": t.col_count, + } + ) + parms.tab_rects = tab_rects # list of table rectangles parms.tab_rects0 = list(tab_rects.values()) diff --git a/src/helpers/table_html/__init__.py b/src/helpers/table_html/__init__.py new file mode 100644 index 00000000..1130c8ca --- /dev/null +++ b/src/helpers/table_html/__init__.py @@ -0,0 +1,15 @@ +"""HTML table reconstruction for pymupdf4llm. + +Public entry points: + +* `page_html_tables(page)` -- the per-page `(bbox, html, rows, cols, cells, + extract)` payload consumed by `pymupdf_rag.to_markdown(..., table_output="html")` + and `document_layout`; +* `to_html(pdf, page_index)` -- a standalone page -> `
` HTML convenience. +""" + +from __future__ import annotations + +from .reconstruct import to_html, page_html_tables + +__all__ = ["to_html", "page_html_tables"] diff --git a/src/helpers/table_html/core.py b/src/helpers/table_html/core.py new file mode 100644 index 00000000..2c7a2f10 --- /dev/null +++ b/src/helpers/table_html/core.py @@ -0,0 +1,10 @@ +from __future__ import annotations +"""Document-level HTML join used by :func:`reconstruct.to_html`.""" +from typing import Any + + +def html_document(tables: list[dict[str, Any]]) -> str: + """Concatenate per-table HTML fragments into one document string. + + Joins each entry's ``"html"`` (skipping empties) with a blank line.""" + return "\n\n".join(str(table.get("html") or "") for table in tables if table.get("html")) diff --git a/src/helpers/table_html/reconstruct.py b/src/helpers/table_html/reconstruct.py new file mode 100644 index 00000000..980752d1 --- /dev/null +++ b/src/helpers/table_html/reconstruct.py @@ -0,0 +1,113 @@ +from __future__ import annotations +"""HTML table reconstruction helpers for pymupdf4llm. + +Builds the per-page table payload consumed by the markdown/JSON renderers +(:func:`page_html_tables`) and a standalone page-to-HTML convenience +(:func:`to_html`), both driven by ``page.find_tables(use_layout=True, +union=True, refine=True)``. +""" +import pymupdf +from .core import html_document + + +def _placement_grid_matrices(placements) -> tuple[int, int, list, list]: + """Derive ``(row_count, col_count, cells, extract)`` from a placement grid. + + ``placements`` is a ``Table.placements`` row-major grid of ``SpanCell`` + colspan/rowspan cells. ``row_count`` is the ```` count; ``col_count`` the + column extent after resolving spans (the HTML column count). ``cells`` and + ``extract`` are the ``row_count x col_count`` post-span bbox and plain-text + matrices, ``None`` where a span covers a slot or the grid has a gap.""" + row_count = len(placements) + # Column extent after resolving colspan/rowspan == the HTML ", re.S) + + +def _table_row_width(row_html): + """Row width in grid columns, summing each cell's colspan (default 1).""" + width = 0 + for attrs in _CELL_TAG_RE.findall(row_html): + m = _COLSPAN_RE.search(attrs) + width += int(m.group(1)) if m else 1 + return width + + +def test_to_json_html_mode_grid_fields_consistent(): + """row_count/col_count/cells/extract reported next to html must describe the + SAME grid the html shows; markdown stays unset (html authoritative).""" + js = pymupdf4llm.to_json( + TABLE_PDF, + pages=[0], + table_output="html", + use_ocr=False, + ) + data = json.loads(js) + table_boxes = [ + box + for page in data["pages"] + for box in page["boxes"] + if box["boxclass"] == "table" + ] + assert table_boxes # guard against an empty (vacuously-passing) run + + for box in table_boxes: + tbl = box["table"] + html = tbl["html"] + rows = _ROW_TAG_RE.findall(html) + + assert len(rows) == tbl["row_count"] + assert max((_table_row_width(row) for row in rows), default=0) == tbl["col_count"] + + # Shape, not full occupancy: individual cells may legitimately be + # None (grid gap / covered by a span), but row/col dimensions must + # match row_count/col_count. + cells = tbl["cells"] + assert len(cells) == tbl["row_count"] + assert all(len(row) == tbl["col_count"] for row in cells) + + # extract is the post-span cell-text matrix (row-major, None for a + # span-covered slot / grid gap), same shape as cells. markdown stays None + # (html authoritative in this mode). + extract = tbl["extract"] + assert extract is not None + assert len(extract) == tbl["row_count"] + assert all(len(row) == tbl["col_count"] for row in extract) + assert tbl["markdown"] is None + + +def test_table_output_html_no_layout_falls_back_to_rag_path(monkeypatch): + """When the layout engine is unavailable, table_output="html" must still + work end-to-end by falling back to the legacy pymupdf_rag path (which has + its own independent table_output="html" wiring).""" + monkeypatch.setattr(pymupdf4llm, "_use_layout", False) + + md = pymupdf4llm.to_markdown(TABLE_PDF, table_output="html") + + assert "
/ column + # count the matrices below are shaped to. + occupied = set() + col_count = 0 + for row_idx, row in enumerate(placements): + col_idx = 0 + for cell in row: + while (row_idx, col_idx) in occupied: + col_idx += 1 + for dr in range(cell.rowspan): + for dc in range(cell.colspan): + occupied.add((row_idx + dr, col_idx + dc)) + col_idx += cell.colspan + col_count = max(col_count, col_idx) + # Expand the ragged post-span grid into row_count x col_count bbox and text + # matrices, ``None`` where a span covers a slot (or a grid gap). + bbox_grid = [[None] * col_count for _ in range(row_count)] + text_grid = [[None] * col_count for _ in range(row_count)] + covered = set() + for row_idx, row in enumerate(placements): + col_idx = 0 + for cell in row: + while (row_idx, col_idx) in covered: + col_idx += 1 + if col_idx >= col_count: + break + bbox_grid[row_idx][col_idx] = list(cell.bbox) if cell.bbox is not None else None + text_grid[row_idx][col_idx] = cell.text + for dr in range(cell.rowspan): + for dc in range(cell.colspan): + if dr or dc: + covered.add((row_idx + dr, col_idx + dc)) + col_idx += cell.colspan + return row_count, col_count, bbox_grid, text_grid + + +def to_html(pdf, page_index=0): + """Reconstruct the tables on one PDF page and return them as an HTML string. + + pdf : PDF file path (str/Path) or an already-open pymupdf.Document. + page_index : 0-based page number. + returns : concatenated ...
HTML (empty string if no tables). + + When ``pdf`` is an already-open Document, the target page is derotated in + place via ``page.remove_rotation()`` (mutates the caller-owned page). + """ + owns_doc = not isinstance(pdf, pymupdf.Document) + doc = pymupdf.open(pdf) if owns_doc else pdf + try: + page = doc[page_index] + page.remove_rotation() + tf = page.find_tables(use_layout=True, union=True, refine=True) + tables = [tab.to_html() for tab in (getattr(tf, "tables", None) or [])] + finally: + if owns_doc: + doc.close() + return html_document([{"html": h} for h in tables]) + + +def page_html_tables(page: pymupdf.Page) -> list[tuple[pymupdf.Rect, str, int, int, list, list]]: + """Reconstruct one already-open page's tables as payload tuples. + + Returns one ``(bbox, html, rows, cols, cells, extract)`` tuple per table, in + reading order, for the markdown/JSON renderers to drive table emission, + reading order and body-text exclusion: + + * ``bbox`` -- ``tab.bbox`` (a grid-ref table keeps its reported layout box); + * ``html`` -- ``tab.to_html()``; + * ``rows``/``cols``/``cells``/``extract`` -- the reconstructed grid the + ``html`` shows (see :func:`_placement_grid_matrices`): ``extract`` is the + per-cell plain-text matrix (``None`` for span-covered slots / grid gaps) + and ``cells`` the matching post-span bbox matrix, both ``rows x cols``. + + The caller must remove page rotation before calling. Not thread-safe on a + shared Page: core caches word/vector extraction as attributes on the given + ``page``, so concurrent calls must each use their own ``pymupdf.Page``. + """ + tf = page.find_tables(use_layout=True, union=True, refine=True) + result = [] + for tab in (getattr(tf, "tables", None) or []): + row_count, col_count, cells, extract = _placement_grid_matrices(tab.placements) + result.append( + ( + pymupdf.Rect(tab.bbox), + tab.to_html(), + row_count, + col_count, + cells, + extract, + ) + ) + return result diff --git a/tests/test_table_html.py b/tests/test_table_html.py new file mode 100644 index 00000000..02c1cbff --- /dev/null +++ b/tests/test_table_html.py @@ -0,0 +1,293 @@ +import inspect +import json +import os +import re +from concurrent.futures import ThreadPoolExecutor + +import pymupdf +import pymupdf4llm +from pymupdf4llm.helpers import document_layout +from pymupdf4llm.helpers.table_html import page_html_tables +from pymupdf4llm.helpers.table_html.reconstruct import to_html + + +g_root = os.path.normpath(f"{__file__}/../..") +TABLE_PDF = os.path.join(g_root, "tests", "test_sce_150_1.pdf") + + +def test_to_html_is_live_only_public_api(): + signature = inspect.signature(to_html) + assert list(signature.parameters) == ["pdf", "page_index"] + + +def test_page_html_tables_uses_core_union_find_tables(): + original_find_tables = pymupdf.Page.find_tables + calls = [] + + def wrapped_find_tables(self, *args, **kwargs): + calls.append(kwargs.copy()) + return original_find_tables(self, *args, **kwargs) + + pymupdf.Page.find_tables = wrapped_find_tables + try: + doc = pymupdf.open(TABLE_PDF) + try: + tables = page_html_tables(doc[0]) + finally: + doc.close() + finally: + pymupdf.Page.find_tables = original_find_tables + + assert len(tables) == 2 + # find_tables(union=True, refine=True) is called exactly once at the Page + # level; the line-based candidate pass runs via the module-level find_tables + # (not a Page.find_tables call), so it is not counted here. + assert len(calls) == 1 + assert calls[0]["use_layout"] is True + assert calls[0]["union"] is True + assert calls[0]["refine"] is True + + +def test_to_markdown_table_output_html_uses_layout_path(): + original_parse_document = document_layout.parse_document + calls = [] + + def wrapped_parse_document(*args, **kwargs): + calls.append(kwargs.copy()) + return original_parse_document(*args, **kwargs) + + document_layout.parse_document = wrapped_parse_document + pymupdf4llm.use_layout(True) + try: + md = pymupdf4llm.to_markdown( + TABLE_PDF, + pages=[0], + table_output="html", + use_ocr=False, + ) + finally: + document_layout.parse_document = original_parse_document + assert md.count("", re.S) + + +def test_to_json_html_tables_match_to_markdown(): + """VALUE parity: every ...
emitted into html-mode markdown + must be byte-identical to (and in the same order as) the html the + same-mode JSON reports for its table boxes.""" + md = pymupdf4llm.to_markdown( + TABLE_PDF, + pages=[0], + table_output="html", + use_ocr=False, + ) + js = pymupdf4llm.to_json( + TABLE_PDF, + pages=[0], + table_output="html", + use_ocr=False, + ) + data = json.loads(js) + table_boxes = [ + box + for page in data["pages"] + for box in page["boxes"] + if box["boxclass"] == "table" + ] + json_html = "\n\n".join( + box["table"]["html"] for box in table_boxes if box["table"].get("html") + ) + + md_tables = _TABLE_TAG_RE.findall(md) + json_tables = _TABLE_TAG_RE.findall(json_html) + + assert md_tables # guard against both-empty passes + assert md_tables == json_tables + + +_CELL_TAG_RE = re.compile(r"<(?:td|th)\b([^>]*)>") +_COLSPAN_RE = re.compile(r'colspan="(\d+)"') +_ROW_TAG_RE = re.compile(r"]*>.*?