Files
2026-07-13 12:08:54 +08:00

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)