pinion-pdf/pinion/convert.py

"""Rebuilding spreadsheet rows from the measured pages.

Walks the batches, gathers each row's cells from its set of pages, files
text by where it STARTS (Trap 1), notes blackout bars (Trap 2), and writes
rows out in order while keeping score for the final checks.
"""

from .page import all_runs, run_start, run_text, row_numbers, solid_regions


def which_column(edges, x):
    """Which column does text starting at x belong to? Checked against the
    START of the text (Trap 1), with a slight fudge for boundary wobble."""
    for i in range(len(edges) - 1):
        if edges[i] - 0.6 <= x < edges[i + 1] - 0.6:
            return i
    return None


def convert(doc, M, out, n_batches):
    """Rebuild every row and write it to `out` (a csv.writer). Returns a
    stats dict for the checks at the end."""
    rpp, pitch = M["rows_per_page"], M["pitch"]
    stats = {"text_no_column": 0, "text_no_row": 0, "bars_no_row": 0,
             "pieces": 0, "gutter_disagreed": 0, "gutter_unreadable": 0,
             "n_rows": 0, "chain_ok": True,
             "blackouts": [], "readable": []}
    next_expected = 1

    for batch in range(n_batches):
        if batch and batch % max(1, n_batches // 10) == 0:
            print(f"  {100 * batch // n_batches}% ...")
        cells, redacted, known_rows = {}, {}, set()

        for slice_no in range(M["ppb"]):
            page = doc[batch * M["ppb"] + slice_no]
            edges, offset = M["bands"][slice_no], M["offsets"][slice_no]
            runs = all_runs(page)
            data = [r for r in runs if r["font"] != M["font"]]
            if not data:
                continue

            # STEP 1 - Read this page's printed row numbers. They tell us
            # exactly which rows it shows, so we never have to assume.
            printed = row_numbers(runs, M["font"], M["table_left"])
            stats["gutter_unreadable"] += rpp - len(printed)

            # STEP 2 - Vertical position -> row number. Anchor on one
            # printed number and step by whole rows from there; this
            # survives redactions that happen to cover a row number.
            if printed:
                anchor_y = min(printed)
                anchor_n = printed[anchor_y]
            else:
                anchor_y = M["first_row_y"]
                anchor_n = batch * rpp + 1

            def row_at(y):
                n = anchor_n + round((y - anchor_y) / pitch)
                lo = anchor_n - (anchor_n - 1) % rpp
                return n if lo <= n < lo + rpp else None

            # Every printed number we CAN read is a free check on that
            # rule. If anything goes off the rails, this will let us know.
            for y, n in printed.items():
                if row_at(y) != n:
                    stats["gutter_disagreed"] += 1
                known_rows.add(n)

            # STEP 3 - Text -> cell.
            for run in data:
                x, y = run_start(run)
                row = row_at(y)
                if row is None:
                    stats["text_no_row"] += 1
                    continue
                col = which_column(edges, x)
                if col is None:
                    stats["text_no_column"] += 1
                    continue
                slot = cells.setdefault(row, [[] for _ in range(M["n_cols"])])
                slot[offset + col].append((x, run_text(run)))

            # STEP 4 - Blackout bars (Trap 2).
            for box in solid_regions(page):
                if box.width <= M["bar_min_w"] or box.height <= M["bar_min_h"]:
                    continue
                row = row_at((box.y0 + box.y1) / 2)
                if row is None:
                    stats["bars_no_row"] += 1
                    continue
                col = which_column(edges, box.x0)
                letters = M["strips"][slice_no]
                label = letters[col] if col is not None else "?"
                stats["blackouts"].append((row, label, round(box.width, 1)))
                redacted[row] = max(redacted.get(row, 0), round(box.width, 1))

                # Did any writing survive underneath the bar?
                # If so, point and laugh. And also report it.
                for run in data:
                    if row_at(run_start(run)[1]) != row:
                        continue
                    under = [chr(c[0]) for c in run["chars"]
                             if box.x0 - 0.5 <= c[2][0] <= box.x1 - 1.5]
                    if under:
                        stats["readable"].append((row, label, "".join(under)))

        # STEP 5 - Write it all out, in row order. known_rows keeps a
        # completely empty spreadsheet row from vanishing: its printed
        # number proves it exists even when no cell has text.
        for row in sorted(set(cells) | known_rows):
            slot = cells.get(row) or [[] for _ in range(M["n_cols"])]
            line = []
            for cell in slot:
                if len(cell) > 1:
                    stats["pieces"] += 1
                # sorting by x rejoins a cell drawn in several pieces
                line.append("".join(t for _, t in sorted(cell)).strip())
            if row == 1:   # the header row names our two extra columns
                out.writerow(line + ["reason_redacted", "redaction_bar_pt"])
            else:
                bar = redacted.get(row)
                out.writerow(line + ["TRUE" if bar else "FALSE", bar or ""])
                stats["n_rows"] += 1
            # Row numbers must count 1,2,3... with no gaps or repeats. Duh.
            if row != next_expected:
                stats["chain_ok"] = False
            next_expected = row + 1

    return stats