Extraire les tableaux d’un PDF

"""
Lire les tableaux d'un PDF, et dire comment ils ont été lus.

Niveau N0. Un tableau dans un PDF n'est pas un tableau : c'est du texte posé
à des coordonnées, parfois avec des traits tracés autour. Il en découle deux
lectures, et elles ne se valent pas.

Quand le tableau est tracé, les filets disent où sont les cellules, et la
lecture est exacte. C'est `pdfplumber` qui le fait ici, et c'est la seule
chose que cette fiche n'écrit pas elle-même : trouver les filets, les
apparier en grille et affecter les mots aux cases est son travail, et elle le
fait bien.

Quand il n'y a pas de filet, il faut deviner les colonnes d'après la position
des mots. Cette part-là est écrite plus bas, dans les deux langages, pour
qu'ils répondent la même chose — il n'existe pas d'équivalent de `pdfplumber`
en JavaScript, et le repli de `pdfplumber` lui-même coupe une cellule plus
large que sa colonne, transformant « Prix HT » en « Prix H ».

Laquelle des deux lectures a répondu figure dans le rapport, parce que
l'appelant ne devrait pas avoir à se le demander : `strategy` vaut « lines »
quand la page disait où étaient les cellules, et « text » quand c'est le code
qui a décidé. Seul le Python peut répondre « lines » : aucun paquet
JavaScript ne lit les filets d'un PDF, donc ce côté-là devine toujours.

Le devinement porte sur toute la page, pas sur le tableau : les colonnes
viennent de tous les bords gauches vus sur la page, bloc adresse et pied de
page compris. C'est pourquoi une lecture devinée ne garde que les lignes à au
moins deux cellules remplies — « SARL Le Moulin » seule sur sa ligne est
l'en-tête de la lettre, pas une ligne du tableau. Ce qu'elle écarte n'est pas
jeté : `dropped_lines` porte le texte de chaque ligne refusée, si bien que
rien de ce qui a été lu sur la page n'est perdu, et que ce qui n'est pas une
ligne est nommé plutôt que rendu comme telle. Sans cette règle, les quatre
lignes d'en-tête d'une facture reçue par courriel ressortent en quatre lignes
à trois cellules vides, et un appelant qui parcourt `rows[1:]` lit une adresse
de rue comme une référence d'article.
"""

from __future__ import annotations

import io

# Two words are on the same line if their baselines are within this many
# points. A ten-point font makes a line about twelve points tall, so three is
# generous without merging neighbouring rows.
SAME_LINE = 3.0

# Two words belong to the same cell when the gap between them is no wider than
# this. A space in a ten-point font is under three points; the gap between two
# columns of a table is tens of points. Ten sits between the two.
SAME_CELL = 10.0

# Two cells are in the same column when their left edges are within this.
SAME_COLUMN = 12.0

# In a guessed reading, a line with fewer filled cells than this is not a row
# of the table: it is the letterhead, the invoice number, a page footer, a
# paragraph of prose, or the second half of a designation that wrapped. All of
# them go to `dropped_lines`, with their text.
MIN_FILLED_CELLS = 2


def read_tables(pdf_bytes, *, pages=None) -> dict:
    """
    Every table of the document, page by page, with the reading that found it.

    `pages` is the list of page numbers to look at, counted from one; left out,
    every page is read.
    """
    try:
        import pdfplumber

        document = pdfplumber.open(io.BytesIO(pdf_bytes))
    except Exception as error:  # noqa: BLE001 - a broken file is an answer too
        return {"tables": [], "reason": f"this file could not be opened as a PDF: {error}"}

    found = []
    with document:
        for number, page in enumerate(document.pages, start=1):
            if pages is not None and number not in pages:
                continue
            ruled = page.extract_tables()
            if ruled:
                found.extend({"page": number, "strategy": "lines", "rows": rows,
                              "dropped_lines": []}
                             for rows in ruled)
                continue
            words = [{"text": w["text"], "x": w["x0"], "end": w["x1"], "y": -w["top"]}
                     for w in page.extract_words()]
            rows = group_into_rows(words)
            kept = [row for row in rows if _filled(row) >= MIN_FILLED_CELLS]
            dropped = [" ".join(c for c in row if c)
                       for row in rows if _filled(row) < MIN_FILLED_CELLS]
            if kept:
                found.append({"page": number, "strategy": "text", "rows": kept,
                              "dropped_lines": dropped})
    return {"tables": found, "reason": None}


def _filled(row) -> int:
    """How many cells of this line carry text."""
    return sum(1 for cell in row if cell)


def group_into_rows(words) -> list:
    """
    Words with coordinates, turned into a grid.

    Lines first, by baseline; then the columns, taken from the left edges seen
    across the whole page. A word is never cut: it belongs to the column its
    left edge is nearest, which is what keeps « Prix HT » whole.
    """
    if not words:
        return []

    # Lines, by baseline.
    lines = []
    for word in sorted(words, key=lambda w: (-w["y"], w["x"])):
        if lines and abs(lines[-1][0]["y"] - word["y"]) <= SAME_LINE:
            lines[-1].append(word)
        else:
            lines.append([word])

    # Cells, by the gap between words: « Moulin a cafe » is one cell, and the
    # eighty points that follow it are a column boundary.
    celled = []
    for line in lines:
        cells = []
        for word in line:
            if cells and word["x"] - cells[-1]["end"] <= SAME_CELL:
                cells[-1] = {"text": f"{cells[-1]['text']} {word['text']}",
                             "x": cells[-1]["x"], "end": word["end"]}
            else:
                cells.append(dict(word))
        celled.append(cells)

    # Columns, from the left edges seen across the whole page.
    columns = []
    for x in sorted(cell["x"] for line in celled for cell in line):
        if not columns or x - columns[-1] > SAME_COLUMN:
            columns.append(x)

    grid = []
    for line in celled:
        row = [""] * len(columns)
        for cell in line:
            index = min(range(len(columns)), key=lambda i: abs(columns[i] - cell["x"]))
            row[index] = f"{row[index]} {cell['text']}".strip()
        grid.append(row)
    return grid

/**
 * Lire les tableaux d'un PDF, et dire comment ils ont été lus.
 *
 * Niveau N0. Un tableau dans un PDF n'est pas un tableau : c'est du texte posé
 * à des coordonnées, parfois avec des traits tracés autour. Il en découle deux
 * lectures, et elles ne se valent pas.
 *
 * Quand le tableau est tracé, les filets disent où sont les cellules, et la
 * lecture est exacte. C'est ce que fait `pdfplumber` du côté Python de cette
 * fiche, et il n'en existe pas d'équivalent en JavaScript : aucun paquet ne lit
 * les filets d'un PDF pour les apparier en grille. Ce fichier ne fait donc que
 * l'autre lecture, celle qui devine les colonnes d'après la position des mots,
 * et `strategy` le dit sur chaque tableau qu'il rend.
 *
 * La conséquence est mesurée dans le test, et c'est la raison pour laquelle
 * cette fiche rend la stratégie employée : sur un tableau tracé dont la
 * désignation se replie sur deux lignes, l'extrait Python garde les deux
 * moitiés dans la même case, et celui-ci sort la seconde du tableau — elle
 * n'a qu'une cellule remplie, donc elle est nommée dans `dropped_lines`.
 *
 * Le regroupement plus bas est le même code dans les deux langages, pour que
 * sur un tableau sans filets — le cas ordinaire d'une facture reçue par
 * courriel — les deux répondent exactement la même chose.
 *
 * Ce devinement porte sur toute la page, pas sur le tableau : les colonnes
 * viennent de tous les bords gauches vus sur la page, bloc adresse et pied de
 * page compris. C'est pourquoi la lecture ne garde que les lignes à au moins
 * deux cellules remplies — « SARL Le Moulin » seule sur sa ligne est l'en-tête
 * de la lettre, pas une ligne du tableau. Ce qu'elle écarte n'est pas jeté :
 * `dropped_lines` porte le texte de chaque ligne refusée, si bien que rien de
 * ce qui a été lu sur la page n'est perdu. Sans cette règle, les quatre lignes
 * d'en-tête d'une facture reçue par courriel ressortent en quatre lignes à
 * trois cellules vides, et un appelant qui parcourt `rows.slice(1)` lit une
 * adresse de rue comme une référence d'article.
 */

// Two words are on the same line if their baselines are within this many
// points. A ten-point font makes a line about twelve points tall, so three is
// generous without merging neighbouring rows.
export const SAME_LINE = 3.0;

// Two words belong to the same cell when the gap between them is no wider than
// this. A space in a ten-point font is under three points; the gap between two
// columns of a table is tens of points. Ten sits between the two.
export const SAME_CELL = 10.0;

// Two cells are in the same column when their left edges are within this.
export const SAME_COLUMN = 12.0;

// A line with fewer filled cells than this is not a row of the table: it is
// the letterhead, the invoice number, a page footer, a paragraph of prose, or
// the second half of a designation that wrapped. All of them go to
// `dropped_lines`, with their text.
export const MIN_FILLED_CELLS = 2;

/** How many cells of this line carry text. */
const filled = (row) => row.filter((cell) => cell).length;

/**
 * Every table of the document, page by page, with the reading that found it.
 *
 * `pages` is the list of page numbers to look at, counted from one; left out,
 * every page is read.
 *
 * @param {Uint8Array} pdfBytes
 * @param {{pages?: number[], load?: object}} [options]
 */
export async function readTables(pdfBytes, { pages, load } = {}) {
  const pdfjs = load ?? (await import('pdfjs-dist/legacy/build/pdf.mjs'));
  const found = [];
  try {
    const document = await pdfjs.getDocument({
      data: new Uint8Array(pdfBytes), isEvalSupported: false,
    }).promise;
    for (let number = 1; number <= document.numPages; number += 1) {
      if (pages && !pages.includes(number)) continue;
      // eslint-disable-next-line no-await-in-loop -- pdf.js hands pages out one at a time
      const content = await (await document.getPage(number)).getTextContent();
      const words = content.items
        .filter((item) => (item.str ?? '').trim())
        .map((item) => ({
          text: item.str.trim(),
          x: item.transform[4],
          end: item.transform[4] + item.width,
          y: item.transform[5],
        }));
      const rows = groupIntoRows(words);
      const kept = rows.filter((row) => filled(row) >= MIN_FILLED_CELLS);
      const dropped = rows.filter((row) => filled(row) < MIN_FILLED_CELLS)
        .map((row) => row.filter((cell) => cell).join(' '));
      if (kept.length > 0) {
        found.push({ page: number, strategy: 'text', rows: kept, dropped_lines: dropped });
      }
    }
  } catch (error) {
    return { tables: [], reason: `this file could not be opened as a PDF: ${error.message}` };
  }
  return { tables: found, reason: null };
}

/**
 * Words with coordinates, turned into a grid.
 *
 * Lines first, by baseline; then the cells, by the gap between words; then the
 * columns, taken from the left edges seen across the whole page. A word is
 * never cut: it belongs to the column its left edge is nearest.
 */
export function groupIntoRows(words) {
  if (words.length === 0) return [];

  // Lines, by baseline.
  const lines = [];
  for (const word of [...words].sort((a, b) => b.y - a.y || a.x - b.x)) {
    const last = lines.at(-1);
    if (last && Math.abs(last[0].y - word.y) <= SAME_LINE) last.push(word);
    else lines.push([word]);
  }

  // Cells, by the gap between words: « Moulin a cafe » is one cell, and the
  // eighty points that follow it are a column boundary.
  const celled = lines.map((line) => {
    const cells = [];
    for (const word of line) {
      const last = cells.at(-1);
      if (last && word.x - last.end <= SAME_CELL) {
        last.text = `${last.text} ${word.text}`;
        last.end = word.end;
      } else {
        cells.push({ ...word });
      }
    }
    return cells;
  });

  // Columns, from the left edges seen across the whole page.
  const columns = [];
  for (const x of celled.flat().map((cell) => cell.x).sort((a, b) => a - b)) {
    if (columns.length === 0 || x - columns.at(-1) > SAME_COLUMN) columns.push(x);
  }

  return celled.map((line) => {
    const row = Array(columns.length).fill('');
    for (const cell of line) {
      let index = 0;
      for (let i = 1; i < columns.length; i += 1) {
        if (Math.abs(columns[i] - cell.x) < Math.abs(columns[index] - cell.x)) index = i;
      }
      row[index] = `${row[index]} ${cell.text}`.trim();
    }
    return row;
  });
}