Files
pdf/tests/edits/forensic_extract.py
T
2026-07-30 16:48:46 +05:30

408 lines
15 KiB
Python

"""Stage 1-4 extraction + frontend layout mirror (ParagraphEditor.tsx)."""
from __future__ import annotations
import hashlib
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "gateway"))
import pdfengine # type: ignore
def line_advances(line) -> tuple[dict[int, list[float]], float]:
runs = line.runs
seq: list[tuple[int, int, float]] = []
aligned: dict[int, bool] = {}
for ri, r in enumerate(runs):
gs = r.glyphs
ok = bool(r.text) and len(gs) == len(r.text)
aligned[ri] = ok
if ok:
for ci, g in enumerate(gs):
seq.append((ri, ci, g.origin_x))
per_run: dict[int, list[float]] = {}
for ri, r in enumerate(runs):
if aligned.get(ri):
per_run[ri] = [0.0] * len(r.text)
anchor_x = seq[0][2] if seq else line.x
for k, (ri, ci, ox) in enumerate(seq):
gs = runs[ri].glyphs
g = gs[ci]
char_w = g.bbox_w if g.bbox_w > 0 else (runs[ri].font_size or 12) * 0.5
per_run[ri][ci] = (seq[k + 1][2] - ox) if k + 1 < len(seq) else char_w
for ri in list(per_run.keys()):
if any(a <= 0 for a in per_run[ri]):
del per_run[ri]
return per_run, anchor_x
def para_eff_size(lines) -> float:
nominal = 0.0
for line in lines:
for r in line.runs:
sz = max(r.font_size or 0, r.h or 0)
if sz > nominal:
nominal = sz
return nominal or 12.0
def median(xs: list[float]) -> float:
if not xs:
return 0.0
s = sorted(xs)
return s[len(s) // 2]
def compute_layout(para) -> dict:
lines = para.lines
eff_size = para_eff_size(lines)
column_left = float("inf")
first_baseline_y = float("-inf")
baselines: list[float] = []
right_edges: list[float] = []
object_indices: list[int] = []
seed_runs: list[dict] = []
orig_lines: list[dict] = []
for li, line in enumerate(lines):
if line.baseline_y is not None:
baselines.append(line.baseline_y)
first_baseline_y = max(first_baseline_y, line.baseline_y)
column_left = min(column_left, line.x)
right_edges.append(line.x + line.w)
line_runs = line.runs
per_run, anchor_x = line_advances(line)
line_frags: list[dict] = []
for ri, r in enumerate(line_runs):
for oi in (r.object_indices or []):
object_indices.append(oi)
orig = r.text or ""
text = orig
if li > 0 and ri == 0 and seed_runs:
prev = seed_runs[-1]["text"]
if prev and not prev.endswith((" ", "\t")) and not text.startswith((" ", "\t")):
text = " " + text
adv = per_run.get(ri)
fc = getattr(r, "fill_color", None) or getattr(r, "color", None) or "#000000"
safe_color = fc if fc and fc != "#ffffff00" else "#000000"
r_size = max(r.font_size or 0, r.h or 0) or eff_size
seed_runs.append({
"text": text, "fid": r.internal_font_id or "", "size": r_size,
"color": safe_color, "fontName": r.font_name or "", "advances": adv,
})
if orig:
line_frags.append({
"text": orig, "fid": r.internal_font_id or "", "size": r_size,
"color": safe_color, "advances": adv,
})
if line_frags:
orig_lines.append({"frags": line_frags, "x": anchor_x, "baselineY": line.baseline_y or 0})
max_right = max(right_edges) if right_edges else column_left + 250
column_right = max(max_right, column_left + 250) if column_left != float("inf") else 250
deltas = [baselines[i] - baselines[i + 1] for i in range(len(baselines) - 1)]
dom_size = next((r["size"] for r in seed_runs if r["text"].strip()), 12.0)
leading = abs(median(deltas)) if deltas else dom_size * 1.2
col_w = column_right - column_left
align = "left"
if len(lines) >= 2:
reaching = sum(1 for i in range(len(right_edges) - 1) if right_edges[i] >= column_right - col_w * 0.04)
if reaching >= (len(lines) - 1) * 0.7:
align = "justify"
return {
"columnLeft": column_left, "columnRight": column_right,
"firstBaselineY": first_baseline_y, "leading": leading,
"oldLineCount": len(lines), "align": align,
"objectIndices": object_indices, "seedRuns": seed_runs, "origLines": orig_lines,
"domSize": dom_size,
}
def extract_flat_runs(seed_runs: list[dict], dominant_fid: str, dom_size: float, dom_color: str) -> list[dict]:
"""Mirror extractFlatRuns when DOM matches seedRuns exactly (edit mode entry)."""
out: list[dict] = []
for r in seed_runs:
if not r["text"]:
continue
frag = {
"text": r["text"], "internalFontId": r["fid"] or dominant_fid,
"fontSize": r["size"], "color": r["color"],
}
adv = r.get("advances")
if adv and len(adv) == len(r["text"]):
frag["advances"] = adv
out.append(frag)
return out
def build_reflow_data(layout: dict, runs: list[dict], orig_lines: list[dict] | None, para_id: str) -> dict:
line_position = {}
if layout["origLines"]:
line_position = {
"lineX": [l["x"] for l in layout["origLines"]],
"lineBaselineY": [l["baselineY"] for l in layout["origLines"]],
}
lines_data = {}
if orig_lines:
lines_data = {
"lines": [
[
{
"text": f["text"], "internalFontId": f["fid"],
"fontSize": f["size"], "color": f["color"],
**({"advances": f["advances"]} if f.get("advances") else {}),
}
for f in l["frags"]
]
for l in orig_lines
],
}
dominant_fid = next((r["fid"] for r in layout["seedRuns"] if r["text"].strip() and r["fid"]), "")
dom_run = next((r for r in layout["seedRuns"] if r["text"].strip() and r["fid"] == dominant_fid), None)
dom_size = (dom_run or layout["seedRuns"][0])["size"] if layout["seedRuns"] else 12.0
dom_color = (dom_run or layout["seedRuns"][0])["color"] if layout["seedRuns"] else "#000000"
out_runs = runs if runs else [{"text": " ", "internalFontId": dominant_fid, "fontSize": dom_size, "color": "#000000"}]
return {
"objectIndices": layout["objectIndices"],
"runs": out_runs,
**line_position,
**lines_data,
"columnLeft": layout["columnLeft"],
"columnRight": layout["columnRight"],
"pushColumnLeft": layout["columnLeft"],
"firstBaselineY": layout["firstBaselineY"],
"leading": layout["leading"],
"oldLineCount": layout["oldLineCount"],
"align": layout["align"],
"paraId": para_id,
}
def paragraph_snapshot_from_extraction(para, layout: dict) -> dict:
"""Original PDF paragraph metrics from extraction + derived advances."""
lines = para.lines
all_text = "".join(r.text for ln in lines for r in ln.runs)
runs_detail = []
for ln in lines:
per_run, anchor_x = line_advances(ln)
for ri, r in enumerate(ln.runs):
adv = per_run.get(ri)
runs_detail.append({
"text": r.text,
"fontId": r.internal_font_id,
"fontName": r.font_name,
"fontSize": max(r.font_size or 0, r.h or 0),
"color": getattr(r, "fill_color", None) or getattr(r, "color", None) or "#000000",
"glyphCount": len(r.glyphs),
"advanceCount": len(adv) if adv else 0,
"advances": adv,
"bbox": [r.x, r.y, r.w, r.h],
"lineWidth": sum(adv) if adv else r.w,
})
line_widths = [ln.w for ln in lines]
orig_line_texts = ["".join(r.text for r in ln.runs) for ln in lines]
return {
"stage": "ORIGINAL_PDF_EXTRACTION",
"text": all_text,
"fontId": runs_detail[0]["fontId"] if runs_detail else "",
"fontName": runs_detail[0]["fontName"] if runs_detail else "",
"fontSize": layout["domSize"],
"color": runs_detail[0]["color"] if runs_detail else "#000000",
"glyphCount": sum(len(r.glyphs) for ln in lines for r in ln.runs),
"advanceCount": sum(rd["advanceCount"] for rd in runs_detail),
"advances": [a for rd in runs_detail for a in (rd["advances"] or [])],
"boundingBox": [para.x, para.y, para.w, para.h],
"paragraphWidth": layout["columnRight"] - layout["columnLeft"],
"lineWidths": line_widths,
"baseline": layout["firstBaselineY"],
"ascent": None,
"descent": None,
"lineHeight": layout["leading"],
"origLines": orig_line_texts,
"runs": runs_detail,
"origLinesProvided": False,
}
def snapshot_from_payload(stage: str, layout: dict, runs: list[dict], orig_lines: list[dict] | None) -> dict:
all_text = "".join(r["text"] for r in runs)
dom = runs[0] if runs else {}
adv_flat = []
for r in runs:
if r.get("advances"):
adv_flat.extend(r["advances"])
return {
"stage": stage,
"text": all_text,
"fontId": dom.get("internalFontId", ""),
"fontName": layout["seedRuns"][0]["fontName"] if layout["seedRuns"] else "",
"fontSize": dom.get("fontSize", layout["domSize"]),
"color": dom.get("color", "#000000"),
"glyphCount": len(all_text),
"advanceCount": len(adv_flat),
"advances": adv_flat,
"boundingBox": None,
"paragraphWidth": layout["columnRight"] - layout["columnLeft"],
"lineWidths": None,
"baseline": layout["firstBaselineY"],
"ascent": None,
"descent": None,
"lineHeight": layout["leading"],
"origLines": ["".join(f["text"] for f in l["frags"]) for l in orig_lines] if orig_lines else None,
"runs": [
{
"text": r["text"],
"fontId": r.get("internalFontId"),
"fontSize": r.get("fontSize"),
"color": r.get("color"),
"advanceCount": len(r.get("advances") or []),
"advances": r.get("advances"),
"advanceSource": "frontend_origin_x_derived" if r.get("advances") else "missing",
}
for r in runs
],
"origLinesProvided": orig_lines is not None,
}
def snapshot_from_wasm_layout(stage: str, layout_json: dict, seed_font_name: str) -> dict:
lines = layout_json.get("lines") or []
all_text = "".join(ln.get("text", "") for ln in lines)
adv_flat = []
line_widths = []
for ln in lines:
adv = ln.get("adv") or []
adv_flat.extend(adv)
line_widths.append(sum(adv))
dom_fs = lines[0].get("fontSize") if lines else None
return {
"stage": stage,
"text": all_text,
"fontId": "",
"fontName": seed_font_name,
"fontSize": dom_fs,
"color": None,
"glyphCount": len(all_text),
"advanceCount": len(adv_flat),
"advances": adv_flat,
"boundingBox": None,
"paragraphWidth": (layout_json.get("columnRight") or 0) - (layout_json.get("columnLeft") or 0),
"lineWidths": line_widths,
"baseline": lines[0].get("baselineY") if lines else None,
"ascent": None,
"descent": None,
"lineHeight": None,
"origLines": [ln.get("text", "") for ln in lines],
"runs": [
{
"text": ln.get("text"),
"lineX0": ln.get("x0"),
"fontSize": ln.get("fontSize"),
"advanceCount": len(ln.get("adv") or []),
"advances": ln.get("adv"),
"advanceSum": sum(ln.get("adv") or []),
"advanceSource": "wasm_reflow_output",
}
for ln in lines
],
"origLinesProvided": None,
}
def para_signature(snap: dict) -> str:
parts = [
snap.get("text") or "",
str(snap.get("fontId") or ""),
str(snap.get("fontSize") or ""),
json.dumps(snap.get("advances") or []),
json.dumps(snap.get("lineWidths") or []),
str(snap.get("baseline") or ""),
json.dumps(snap.get("boundingBox") or []),
]
return hashlib.sha256("".join(parts).encode()).hexdigest()
def main():
pdf_path = ROOT / "corpus" / "basic" / "hello_world.pdf"
if len(sys.argv) > 1:
pdf_path = Path(sys.argv[1])
doc = pdfengine.PdfDocument.load_from_file(str(pdf_path), "")
model = doc.get_page(0).extract_document_model()
para = model.paragraphs[0]
layout = compute_layout(para)
fonts = doc.get_fonts(0, 0)
dominant_fid = next((r["fid"] for r in layout["seedRuns"] if r["text"].strip() and r["fid"]), "")
dom_run = next((r for r in layout["seedRuns"] if r["text"].strip()), layout["seedRuns"][0])
dom_size = dom_run["size"]
dom_color = dom_run["color"]
para_id = f"forensic-{hashlib.md5(str(pdf_path).encode()).hexdigest()[:8]}"
# Stage snapshots
original = paragraph_snapshot_from_extraction(para, layout)
flat_entry = extract_flat_runs(layout["seedRuns"], dominant_fid, dom_size, dom_color)
payload_entry = build_reflow_data(layout, flat_entry, layout["origLines"], para_id)
edit_entry = snapshot_from_payload("EDIT_MODE_ENTRY_PAYLOAD", layout, flat_entry, layout["origLines"])
# After one character
flat_one = []
for i, r in enumerate(flat_entry):
if i == 0:
flat_one.append({**r, "text": r["text"] + "x"})
else:
flat_one.append(r)
payload_one = build_reflow_data(layout, flat_one, None, para_id)
edit_one = snapshot_from_payload("AFTER_ONE_CHAR_PAYLOAD", layout, flat_one, None)
out = {
"pdf": str(pdf_path),
"fonts": [
{
"internalFontId": f.internal_font_id,
"fontName": f.font_name,
"ascent": f.ascent,
"descent": f.descent,
"capHeight": getattr(f, "cap_height", None),
"unitsPerEm": getattr(f, "units_per_em", None),
}
for f in fonts
],
"layout": layout,
"original": original,
"editEntry": edit_entry,
"editOne": edit_one,
"payloads": {
"editEntry": {
"version": "1.0",
"operations": [{
"id": "forensic-entry", "type": "reflow_paragraph", "pageIndex": 0,
"data": payload_entry,
}],
},
"editOne": {
"version": "1.0",
"operations": [{
"id": "forensic-one", "type": "reflow_paragraph", "pageIndex": 0,
"data": payload_one,
}],
},
},
"signatures": {
"original": para_signature(original),
"editEntryPayload": para_signature(edit_entry),
"editOnePayload": para_signature(edit_one),
},
}
out_path = Path(__file__).resolve().parent / "forensic_extract_out.json"
out_path.write_text(json.dumps(out, indent=2), encoding="utf-8")
print(json.dumps(out, indent=2))
if __name__ == "__main__":
main()