Coverage for src / quber / playground / export.py: 96%
107 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""What a batch exports: one row per question, as a workbook or as text.
3The two writers share the row building and nothing else, because each writes a
4different thing. A workbook cell carries its own type and number format, so the
5printed figure is a text cell nothing reformats and a percentage arrives as the
6fraction under a percent format, displaying at the printed scale. A text file
7encodes what it can in the characters: the printed column is quoted, the
8numeric column is not, and a percentage keeps its percent sign for the
9importer to divide. Both land the same value in the same cell; they need
10different bytes to get there.
12Every figure is written twice, as printed and as a number. Both are read off
13the answer — `value` and the derived `number` — and nothing here parses a
14printed figure. `answers/figures.py` is the one place that does, and the two
15sign-flipping defects that motivated that rule were both written by an export
16re-deriving what the answer already carried.
18A question with no answer is a third kind of row, distinct from a figure and
19from a failure: a sentinel in the printed column, an empty numeric column, and
20the model's explanation in the note.
22Formula injection is handled structurally. Workbook text cells are typed as
23text, so a question beginning with `=` is characters. In the CSV the free-text
24columns are prefixed, and only those — a printed figure is never prefixed,
25because that column has to match the filing byte for byte.
26"""
28from __future__ import annotations
30import json
31from decimal import Decimal
32from io import BytesIO
33from typing import Any, Dict, List, Optional
35import openpyxl
36from pydantic import BaseModel
38from quber.playground.answers.figures import is_percent, norm_printed
39from quber.playground.batch import RunState
41COLUMNS = [
42 "question",
43 "as printed",
44 "number",
45 "unit",
46 "period",
47 "page",
48 "cited ids",
49 "source cell text",
50 "status",
51 "note",
52 "review flag",
53 "error",
54]
56#: The printed-column sentinel for an `unanswerable` answer: retrieval ran,
57#: the model answered, and the answer is that the document does not print the
58#: figure. A reader sorting the sheet can then separate a figure, a figure the
59#: document does not contain, and a question that never got that far. The
60#: hashes mark it as ours: it is written off the answer's `unanswerable`
61#: payload kind, never taken from model text, so nothing a model writes and
62#: nothing a filing prints can collide with it.
63NOT_IN_DOCUMENT = "#not found#"
65#: What the review-flag column holds when a cited cell's status is registered
66#: for review. A word rather than a boolean so the filter reads as what it is.
67REVIEW = "review"
69_JOIN = "; "
72class ExportRow(BaseModel):
73 """One question's row, typed for the writers rather than for a reader."""
75 question: str
76 printed: str = ""
77 number: Optional[Decimal] = None
78 percent: bool = False
79 # True when `printed` is a figure quoted verbatim from the filing — the
80 # one kind of cell the CSV writer must never prefix.
81 figure: bool = False
82 unit: str = ""
83 period: str = ""
84 pages: str = ""
85 cited_ids: str = ""
86 cell_text: str = ""
87 status: str = ""
88 note: str = ""
89 review: bool = False
90 error: str = ""
93def build_rows(state: RunState) -> List[ExportRow]:
94 return [_build_row(row.question, row.error, row.response) for row in state.rows]
97def _build_row(question: str, error: Optional[str], response: Optional[Dict[str, Any]]) -> ExportRow:
98 # A question that failed carries its error and nothing else.
99 if response is None:
100 return ExportRow(question=question, error=error or "failed")
102 row = ExportRow(question=question)
103 refs = response.get("references") or []
104 payload = response.get("value") or {}
105 kind = response.get("shape")
107 if kind == "scalar":
108 row.printed = payload.get("value") or ""
109 row.figure = True
110 row.percent = is_percent(row.printed)
111 # `number` was derived by the answer; this is deserialization, not a
112 # second reading of the printed figure.
113 serialized = payload.get("number")
114 row.number = Decimal(serialized) if serialized is not None else None
115 row.unit = payload.get("unit") or ""
116 row.period = payload.get("period") or ""
117 elif kind == "unanswerable":
118 row.printed = NOT_IN_DOCUMENT
119 row.note = payload.get("reason") or ""
120 elif kind == "prose":
121 row.printed = payload.get("text") or ""
122 else:
123 # A series or grid has no single printed figure; the payload itself is
124 # the answer, carried whole rather than truncated to a first point.
125 row.printed = json.dumps(payload, separators=(",", ":"))
126 row.unit = payload.get("unit") or ""
128 seen_pages: List[str] = []
129 for ref in refs:
130 page = str(ref.get("page"))
131 if page not in seen_pages:
132 seen_pages.append(page)
133 row.pages = _JOIN.join(seen_pages)
134 row.cited_ids = _JOIN.join(r.get("ref_id", "") for r in refs)
135 # Where a cited answer rests on more than one cell, each column carries
136 # all of them rather than the first.
137 row.cell_text = _JOIN.join(r["text"] for r in refs if r.get("text"))
138 row.status = _JOIN.join(r["status"] for r in refs if r.get("status"))
139 notes = [r["note"] for r in refs if r.get("note")]
140 if notes:
141 row.note = _JOIN.join(([row.note] if row.note else []) + notes)
142 row.review = any(r.get("flagged") for r in refs)
143 return row
146def percent_format(printed: str) -> str:
147 """A percent format with as many decimals as the figure was printed with,
148 so `6.7 %` shows as 6.7% rather than being padded to 6.70%."""
149 digits = norm_printed(printed).rstrip("%").strip()
150 places = len(digits.split(".")[1]) if "." in digits else 0
151 return f"0.{'0' * places}%" if places else "0%"
154def write_xlsx(state: RunState) -> bytes:
155 """The workbook. Text cells are typed as text — nothing reformats the
156 printed figure and a leading `=` is characters. A percentage's numeric
157 cell holds the fraction under a percent format, so it computes as a rate
158 and displays as the document printed it; the writer picks a value and a
159 format and performs no arithmetic beyond that rescale."""
160 wb = openpyxl.Workbook()
161 ws = wb.active
162 assert ws is not None
163 ws.title = "batch"
164 ws.append(COLUMNS)
165 for r, row in enumerate(build_rows(state), start=2):
166 texts = {
167 1: row.question,
168 2: row.printed,
169 4: row.unit,
170 5: row.period,
171 6: row.pages,
172 7: row.cited_ids,
173 8: row.cell_text,
174 9: row.status,
175 10: row.note,
176 11: REVIEW if row.review else "",
177 12: row.error,
178 }
179 for col, value in texts.items():
180 if value == "":
181 # A blank cell carries nothing to protect; forcing an empty
182 # string into a text cell reads back as an inline-string None.
183 continue
184 cell = ws.cell(row=r, column=col, value=value)
185 cell.data_type = "s"
186 cell.number_format = "@"
187 if row.number is not None:
188 num = ws.cell(row=r, column=3, value=row.number / 100 if row.percent else row.number)
189 if row.percent:
190 num.number_format = percent_format(row.printed)
191 out = BytesIO()
192 wb.save(out)
193 return out.getvalue()
196def _quote(s: str) -> str:
197 return '"' + s.replace('"', '""') + '"'
200def _free_text(s: str) -> str:
201 """A free-text cell: prefixed when it begins with a formula character,
202 then quoted. The prefix is the injection guard; the quotes are what keeps
203 the text as typed when the importer treats quoted fields as text."""
204 if s[:1] in ("=", "+", "-", "@"):
205 s = "'" + s
206 return _quote(s)
209def write_csv(state: RunState) -> str:
210 """The same rows as text. The printed column is quoted and never prefixed;
211 the numeric column is bare, and a percentage keeps its percent sign so the
212 importer stores the fraction under a percent format — the same cell the
213 workbook writes directly."""
214 lines = [",".join(COLUMNS)]
215 for row in build_rows(state):
216 if row.number is None:
217 number = ""
218 elif row.percent:
219 number = f"{row.number}%"
220 else:
221 number = str(row.number)
222 lines.append(
223 ",".join(
224 [
225 _free_text(row.question),
226 _quote(row.printed) if row.figure else _free_text(row.printed),
227 number,
228 _free_text(row.unit),
229 _free_text(row.period),
230 _quote(row.pages),
231 _quote(row.cited_ids),
232 _free_text(row.cell_text),
233 _quote(row.status),
234 _free_text(row.note),
235 REVIEW if row.review else "",
236 _free_text(row.error),
237 ]
238 )
239 )
240 return "\r\n".join(lines) + "\r\n"