pinion-pdf/pinion/convert.py

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 text start with a slight fudge factor."""
    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`."""
    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.
            printed = row_numbers(runs, M["font"], M["table_left"])
            stats["gutter_unreadable"] += rpp - len(printed)

            # STEP 2 - Vertical position -> row number. 
            # This makes sure redaction-covered row numbers don't mess us up.
            # (Yes, it happens)
            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

            # If we can read the number, might as well validate it.
            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.
        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:
                # header row names our two extra redaction 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