192 lines
7.4 KiB
Python
192 lines
7.4 KiB
Python
"""Shared helpers for parsing and re-emitting ``<table>`` markup.
|
|
|
|
These primitives are used by the paragraph-semantic chunker (TableRowSplit
|
|
oversized-table re-split) and by the native multimodal surrounding-context
|
|
extractor. Both call sites need to:
|
|
|
|
* recognise a post-rewrite ``<table id="…" format="…">…</table>`` tag,
|
|
* decide whether the body is JSON or HTML,
|
|
* enumerate row-level units (JSON list items or HTML ``<tr>`` rows along
|
|
with their ``<thead>`` / ``<tbody>`` / ``<tfoot>`` wrappers), and
|
|
* re-serialise a subset of rows while preserving the structural wrappers.
|
|
|
|
Keeping the regexes and helpers in one place avoids subtle drift when
|
|
either consumer evolves.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import html as _html
|
|
import json
|
|
import re
|
|
from typing import Any
|
|
|
|
# Strict regex for a post-rewrite table tag emitted by the sidecar
|
|
# writer (``lightrag.sidecar.writer``):
|
|
# <table id="tb-…" format="json"[ caption="…"]>{rows_json}</table>
|
|
# blocks.jsonl invariants guarantee the tag has no embedded newlines.
|
|
TABLE_TAG_RE = re.compile(
|
|
r"<table\s+(?P<attrs>[^>]*)>(?P<body>.*?)</table>",
|
|
re.DOTALL,
|
|
)
|
|
|
|
# Format detection regex inside the attrs string, e.g. format="json".
|
|
_TABLE_FORMAT_RE = re.compile(r"""format\s*=\s*["'](?P<fmt>[^"']+)["']""")
|
|
|
|
# ``id`` attribute extractor inside the attrs string, e.g. id="tb-…". Used to
|
|
# trace a (possibly row-split) ``<table>`` fragment back to its tables.json
|
|
# entry so consumers can recover per-table metadata (e.g. the repeating header).
|
|
# The negative lookbehind rejects ``id`` that is only the tail of another
|
|
# attribute name — both word-char prefixes (``grid``, ``valid``) and
|
|
# hyphen-joined ones (``data-id``) — so only a standalone ``id`` attribute hits.
|
|
_TABLE_ID_RE = re.compile(r"""(?<![\w-])id\s*=\s*["'](?P<id>[^"']+)["']""")
|
|
|
|
# HTML <tr>...</tr> row extractor. Standard HTML disallows nested <tr>,
|
|
# so a non-greedy match is sufficient for well-formed input.
|
|
HTML_TR_RE = re.compile(r"<tr\b[^>]*>.*?</tr>", re.DOTALL | re.IGNORECASE)
|
|
|
|
# Combined scanner for row-grouping wrappers and rows themselves. Used
|
|
# to attribute each <tr> to its surrounding <thead>/<tbody>/<tfoot> so
|
|
# the wrapper can be reconstructed around chunk boundaries instead of
|
|
# being silently dropped during row-level table splitting.
|
|
HTML_ROW_PARTS_RE = re.compile(
|
|
r"(?P<wrap></?(?:thead|tbody|tfoot)\b[^>]*>)" r"|(?P<tr><tr\b[^>]*>.*?</tr>)",
|
|
re.DOTALL | re.IGNORECASE,
|
|
)
|
|
HTML_WRAPPER_TAG_RE = re.compile(
|
|
r"<(?P<slash>/?)(?P<name>thead|tbody|tfoot)\b", re.IGNORECASE
|
|
)
|
|
|
|
|
|
def detect_table_format(attrs: str, body: str) -> str | None:
|
|
"""Return ``"json"``, ``"html"`` or ``None`` for a parsed ``<table>`` tag.
|
|
|
|
Prefers an explicit ``format="…"`` attribute. When silent, sniffs
|
|
the body: a leading ``[`` / ``{`` (after whitespace) implies JSON;
|
|
the presence of any ``<tr`` tag implies HTML. Anything else is
|
|
unknown and the caller should fall back to character splitting.
|
|
"""
|
|
match = _TABLE_FORMAT_RE.search(attrs or "")
|
|
if match:
|
|
fmt = match.group("fmt").strip().lower()
|
|
if fmt in {"json", "html"}:
|
|
return fmt
|
|
return None
|
|
stripped = (body or "").lstrip()
|
|
if stripped.startswith(("[", "{")):
|
|
return "json"
|
|
if "<tr" in stripped.lower():
|
|
return "html"
|
|
return None
|
|
|
|
|
|
def header_grid_to_thead_html(grid: list[list[str]]) -> str:
|
|
"""Render a header grid (a 2-D string array) as an HTML ``<thead>``.
|
|
|
|
Each row becomes a ``<tr>`` of HTML-escaped ``<th>`` cells inside a single
|
|
``<thead>``. Used by the sidecar writer to serialise an HTML table whose
|
|
recovered header is only available as a grid (e.g. a docx HTML-fallback
|
|
table) — note this cannot reconstruct ``rowspan``/``colspan`` that the grid
|
|
already flattened. Returns ``""`` when no usable row is present.
|
|
"""
|
|
trs: list[str] = []
|
|
for row in grid:
|
|
if not isinstance(row, list):
|
|
continue
|
|
cells = "".join(f"<th>{_html.escape(str(cell))}</th>" for cell in row)
|
|
trs.append(f"<tr>{cells}</tr>")
|
|
return f"<thead>{''.join(trs)}</thead>" if trs else ""
|
|
|
|
|
|
def extract_table_id(attrs: str) -> str | None:
|
|
"""Return the ``id`` attribute value from a ``<table>`` attrs string.
|
|
|
|
``None`` when the attrs carry no ``id`` (e.g. a header already stripped of
|
|
identifiers, or a malformed tag).
|
|
"""
|
|
match = _TABLE_ID_RE.search(attrs or "")
|
|
if match:
|
|
return match.group("id").strip() or None
|
|
return None
|
|
|
|
|
|
def parse_table_tag(text: str) -> tuple[str, list[Any]] | None:
|
|
"""Parse a JSON ``<table …>{rows_json}</table>``.
|
|
|
|
Returns ``(attrs_str, rows)`` or ``None`` if the tag is malformed
|
|
(does not match ``TABLE_TAG_RE``, body is not JSON, or body decodes
|
|
to something other than a list).
|
|
"""
|
|
match = TABLE_TAG_RE.match((text or "").strip())
|
|
if not match:
|
|
return None
|
|
body = match.group("body")
|
|
try:
|
|
rows = json.loads(body)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
if not isinstance(rows, list):
|
|
return None
|
|
return match.group("attrs"), rows
|
|
|
|
|
|
def split_html_rows(body: str) -> list[tuple[str, str]] | None:
|
|
"""Extract ``<tr>...</tr>`` rows tagged with their wrapper context.
|
|
|
|
Returns a list of ``(wrapper_name, tr_str)`` tuples where
|
|
``wrapper_name`` is ``"thead"`` / ``"tbody"`` / ``"tfoot"`` (lower-
|
|
cased) for rows that sit inside the corresponding wrapper, or ``""``
|
|
for rows outside any of those wrappers. ``None`` signals "no row
|
|
found" so the caller falls through to character splitting.
|
|
|
|
Whitespace, captions, comments, ``<colgroup>`` and any other text
|
|
outside the recognised row-wrappers is dropped — this is a regex
|
|
extractor, not a full DOM parser. Wrapper attributes (e.g.
|
|
``<thead class="…">``) are also dropped on re-emission; chunked
|
|
output uses bare wrapper tags.
|
|
"""
|
|
rows: list[tuple[str, str]] = []
|
|
current_wrapper = ""
|
|
for match in HTML_ROW_PARTS_RE.finditer(body or ""):
|
|
wrap = match.group("wrap")
|
|
tr = match.group("tr")
|
|
if wrap is not None:
|
|
tag = HTML_WRAPPER_TAG_RE.match(wrap)
|
|
if tag:
|
|
slash = tag.group("slash")
|
|
name = tag.group("name").lower()
|
|
if slash == "/":
|
|
if current_wrapper == name:
|
|
current_wrapper = ""
|
|
else:
|
|
current_wrapper = name
|
|
elif tr is not None:
|
|
rows.append((current_wrapper, tr))
|
|
if not rows:
|
|
return None
|
|
return rows
|
|
|
|
|
|
def serialize_html_rows(rows: list[tuple[str, str]]) -> str:
|
|
"""Re-emit ``(wrapper, tr)`` rows grouped under their original
|
|
``<thead>`` / ``<tbody>`` / ``<tfoot>`` wrappers.
|
|
|
|
Consecutive rows sharing the same wrapper name collapse into a
|
|
single wrapper block; transitions emit a closing tag for the
|
|
previous wrapper and an opening tag for the next. Rows tagged with
|
|
``""`` (no wrapper) emit bare ``<tr>...</tr>``.
|
|
"""
|
|
parts: list[str] = []
|
|
current_wrapper = ""
|
|
for wrapper, tr in rows:
|
|
if wrapper != current_wrapper:
|
|
if current_wrapper:
|
|
parts.append(f"</{current_wrapper}>")
|
|
if wrapper:
|
|
parts.append(f"<{wrapper}>")
|
|
current_wrapper = wrapper
|
|
parts.append(tr)
|
|
if current_wrapper:
|
|
parts.append(f"</{current_wrapper}>")
|
|
return "".join(parts)
|