Coverage for src / quber / core / figures / grid.py: 98%
102 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Read a scanned page's tables back as grids.
3A scan returns a table as a table. The chunk carries an HTML grid with an id on
4every cell, and the response's grounding map carries, against each of those ids,
5the cell's box and its row, column and spans. Putting the two together gives a
6dense grid — one entry per row and column — with a box on every cell, which is
7the shape the rest of the table pipeline already works in.
9A table is not always its own chunk. Where a page prints a chart and a table as
10one composite, the scan returns one figure, and writes the table inside that
11figure's text as markdown between a pair of markers it emits itself. The values
12are complete and correctly separated there — on one page the scan read a bond
13issuance table's four credit-rating rows exactly, while the parse ran four
14ratings together into a single cell and left the next three blank.
16So the markers are read and the markdown between them is taken as a table.
17Nothing infers where the table is in the text; the scan labelled it.
19What such a table does not carry is geometry. The composite is grounded as one
20box covering the whole figure, with nothing per cell, so its cells arrive with no
21boxes. That costs nothing downstream: the structure correction reads the page
22image rather than the boxes, and a table with no reading to check against already
23puts every cell to the agent.
25Nothing here interprets the content. Spans are honoured as geometry: a cell that
26spans two columns is written once at its own position and the positions it covers
27are left blank, the same way a grid extracted from a text layer arrives. What the
28header block is, which cells merged, and what the title says are read later, off
29the table image, by the steps that read every other table.
30"""
32from __future__ import annotations
34import re
35from html.parser import HTMLParser
36from typing import Any, Dict, List, Optional, Tuple
38from quber.core.extractors.set_of_mark.merge_grounding import markdown_rows
39from quber.core.figures.models import Box, ScannedTable
41#: The grounding entry type carrying one table cell's box and grid position.
42CELL_TYPE = "tableCell"
44#: The grounding entry type carrying the box around a whole table. It is drawn
45#: tighter than the chunk's own box, which reaches out to the surrounding text.
46TABLE_TYPE = "table"
48_TABLE_ID_RE = re.compile(r"<table[^>]*\bid=[\"']([^\"']+)[\"']")
50#: The markers the scan writes around a markdown table it printed inside a
51#: figure's text, when the page prints a chart and a table as one composite.
52_FIGURE_TABLE_RE = re.compile(r"<::table::>(.*?)<::/table::>", re.DOTALL)
55def scanned_tables(response: Dict[str, Any], page: int, job_id: Optional[str]) -> List[ScannedTable]:
56 """Every table in one page's raw response, as a dense grid with per-cell boxes.
58 `page` is the 1-based source page the submission came from. The response
59 numbers its own single page from zero, so the source page is carried in
60 rather than read back out of it.
61 """
62 grounding = response.get("grounding") or {}
63 cells_by_chunk = _cells_by_chunk(grounding)
65 tables: List[ScannedTable] = []
66 for chunk in response.get("chunks") or []:
67 if chunk.get("type") != "table":
68 continue
69 chunk_id = chunk.get("id")
70 markdown = chunk.get("markdown") or ""
71 text_by_id = _cell_text(markdown)
72 cells, boxes = _dense_grid(cells_by_chunk.get(chunk_id, []), text_by_id)
73 tables.append(
74 ScannedTable(
75 page=page,
76 cells=cells,
77 cell_boxes=boxes,
78 box=_table_box(markdown, grounding) or (chunk.get("grounding") or {}).get("box"),
79 chunk_id=chunk_id,
80 job_id=job_id,
81 )
82 )
84 tables.extend(_composite_tables(response, page, job_id))
85 return tables
88def _composite_tables(response: Dict[str, Any], page: int, job_id: Optional[str]) -> List[ScannedTable]:
89 """The tables the scan printed inside a figure's text rather than as chunks.
91 The figure's own box stands as the table's box. It is the only geometry the
92 scan grounded for the composite, and it is drawn around the chart and the
93 table together, so the region is generous rather than wrong.
94 """
95 out: List[ScannedTable] = []
96 for chunk in response.get("chunks") or []:
97 if chunk.get("type") != "figure":
98 continue
99 for markdown in _FIGURE_TABLE_RE.findall(chunk.get("markdown") or ""):
100 rows = markdown_rows(markdown)
101 if not rows:
102 continue
103 width = max(len(r) for r in rows)
104 cells = [r + [""] * (width - len(r)) for r in rows]
105 out.append(
106 ScannedTable(
107 page=page,
108 cells=cells,
109 cell_boxes=[[None] * width for _ in cells],
110 box=(chunk.get("grounding") or {}).get("box"),
111 chunk_id=chunk.get("id"),
112 job_id=job_id,
113 )
114 )
115 return out
118def _cells_by_chunk(grounding: Dict[str, Any]) -> Dict[str, List[Tuple[str, Dict[str, Any]]]]:
119 """The grounding map's cell entries, grouped by the table chunk each belongs to."""
120 by_chunk: Dict[str, List[Tuple[str, Dict[str, Any]]]] = {}
121 for cell_id, entry in grounding.items():
122 if entry.get("type") != CELL_TYPE:
123 continue
124 position = entry.get("position") or {}
125 chunk_id = position.get("chunk_id")
126 if chunk_id is None:
127 continue
128 by_chunk.setdefault(chunk_id, []).append((cell_id, entry))
129 return by_chunk
132def _table_box(markdown: str, grounding: Dict[str, Any]) -> Optional[Box]:
133 """The box around the table itself, from the id its opening tag carries."""
134 match = _TABLE_ID_RE.search(markdown)
135 if match is None:
136 return None
137 entry = grounding.get(match.group(1)) or {}
138 return entry.get("box") if entry.get("type") == TABLE_TYPE else None
141def _dense_grid(
142 cells: List[Tuple[str, Dict[str, Any]]],
143 text_by_id: Dict[str, str],
144) -> Tuple[List[List[str]], List[List[Optional[Box]]]]:
145 """One row per row and one column per column, with a box beside every cell.
147 A cell that spans is written at its own row and column; the positions it
148 covers stay blank and carry no box of their own, so the grid never states the
149 same value twice.
150 """
151 placed: List[Tuple[int, int, str, Optional[Box]]] = []
152 rows = 0
153 columns = 0
154 for cell_id, entry in cells:
155 position = entry.get("position") or {}
156 row = int(position.get("row", 0))
157 col = int(position.get("col", 0))
158 rows = max(rows, row + int(position.get("rowspan", 1) or 1))
159 columns = max(columns, col + int(position.get("colspan", 1) or 1))
160 placed.append((row, col, text_by_id.get(cell_id, ""), entry.get("box")))
162 grid = [["" for _ in range(columns)] for _ in range(rows)]
163 boxes: List[List[Optional[Box]]] = [[None for _ in range(columns)] for _ in range(rows)]
164 for row, col, text, box in placed:
165 if 0 <= row < rows and 0 <= col < columns:
166 grid[row][col] = text
167 boxes[row][col] = box
168 return grid, boxes
171class _CellText(HTMLParser):
172 """Collects the text of every `td`/`th` carrying an id, keyed by that id.
174 Markup inside a cell is dropped and its text kept, so a value typeset with an
175 emphasis or a line break reads as the one string the page prints.
176 """
178 def __init__(self) -> None:
179 super().__init__(convert_charrefs=True)
180 self.text: Dict[str, str] = {}
181 self._open: List[Optional[str]] = []
182 self._parts: List[str] = []
184 def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]) -> None:
185 if tag in ("td", "th"):
186 self._flush()
187 self._open.append(dict(attrs).get("id"))
188 elif tag == "br" and self._open:
189 self._parts.append(" ")
191 def handle_endtag(self, tag: str) -> None:
192 if tag in ("td", "th"):
193 self._flush()
195 def handle_data(self, data: str) -> None:
196 if self._open:
197 self._parts.append(data)
199 def _flush(self) -> None:
200 if not self._open:
201 return
202 cell_id = self._open.pop()
203 if cell_id is not None:
204 self.text[cell_id] = " ".join("".join(self._parts).split())
205 self._parts = []
208def _cell_text(markdown: str) -> Dict[str, str]:
209 """Each cell's printed text, keyed by the id its tag carries."""
210 parser = _CellText()
211 parser.feed(markdown)
212 parser.close()
213 return parser.text