import argparse from contextlib import closing from dataclasses import dataclass from datetime import date from hashlib import sha256 from pathlib import Path import re import sqlite3 from textwrap import wrap from pdfplumber.utils.exceptions import PdfminerException from pypdf.errors import PdfReadError from handbook_search import load_pdf_files from pdf_pages import extract_pages from policy_table import HEADERS, read_edition KINDS = ("prose", "ocr", "table") DEFAULT_DB = Path("mixed-library.sqlite3") @dataclass(frozen=True) class Evidence: source: str page: int kind: str text: str pdf_sha256: str ocr_sha256: str | None effective: str | None equipment_id: str | None table_number: int | None row_number: int | None def citation(self): location = f"{self.source}, page {self.page} [{self.kind}]" if self.kind == "ocr": location += f" | OCR SHA-256 {self.ocr_sha256[:12]}" if self.effective: location += f" | table {self.table_number}, row {self.row_number} | effective {self.effective}" return f"{location} | PDF SHA-256 {self.pdf_sha256[:12]}" def collect(folder): files = sorted(path for path in folder.glob("*.pdf") if path.is_file() and not path.is_symlink()) editions = set() for path in files: if path.name.startswith("equipment-policy-"): edition = read_edition(path) if edition.effective in editions: raise ValueError(f"Multiple policy editions are effective {edition.effective}.") editions.add(edition.effective) for equipment_id, row in edition.rows.items(): text = " | ".join(f"{column}: {value}" for column, value in zip(HEADERS, row.cells)) yield Evidence(path.name, row.page, "table", text, edition.digest, None, edition.effective.isoformat(), equipment_id, row.table, row.row) continue data, sidecars, _ = load_pdf_files(path) pages = extract_pages(data) if sidecars.keys() - {number for number, _ in pages}: raise ValueError(f"{path.name}: an OCR file points outside the PDF.") digest = sha256(data).hexdigest() for number, embedded in pages: sidecar = sidecars.get(number) if embedded and sidecar is not None: raise ValueError(f"{path.name}, page {number}: keep one source of text for this exercise.") kind, ocr_digest = "prose", None if sidecar is not None: embedded = sidecar.decode("utf-8").strip() kind, ocr_digest = "ocr", sha256(sidecar).hexdigest() for text in wrap(" ".join(embedded.split()), width=480): yield Evidence(path.name, number, kind, text, digest, ocr_digest, None, None, None, None) def build(folder, database): data = list(collect(folder)) if not data: raise ValueError(f"No indexable PDFs in {folder}.") with closing(sqlite3.connect(database)) as db, db: db.execute("BEGIN IMMEDIATE") db.execute(""" CREATE TABLE IF NOT EXISTS evidence ( id INTEGER PRIMARY KEY, source TEXT, page INTEGER, kind TEXT, text TEXT, pdf_sha256 TEXT, ocr_sha256 TEXT, effective TEXT, equipment_id TEXT, table_number INTEGER, row_number INTEGER ) """) db.execute("CREATE VIRTUAL TABLE IF NOT EXISTS evidence_words USING fts5(text)") db.execute("DELETE FROM evidence_words") db.execute("DELETE FROM evidence") for item in data: row = db.execute( "INSERT INTO evidence VALUES (NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", (item.source, item.page, item.kind, item.text, item.pdf_sha256, item.ocr_sha256, item.effective, item.equipment_id, item.table_number, item.row_number), ) db.execute("INSERT INTO evidence_words(rowid, text) VALUES (?, ?)", (row.lastrowid, item.text)) for kind in KINDS: print(f"{kind}: {sum(item.kind == kind for item in data)} evidence records") def retrieve(database, as_of, query="", kind=None, equipment=None, limit=3): terms = " OR ".join(f'"{word}"' for word in re.findall(r"[^\W_]+", query)) if not terms and not equipment: raise ValueError("Enter a query with a word, or provide an equipment ID.") if not database.is_file(): raise ValueError("Build the mixed index first.") conditions = ["(e.effective IS NULL OR e.effective = " "(SELECT max(effective) FROM evidence WHERE effective <= ?))"] parameters = [as_of.isoformat()] if kind: conditions.append("e.kind = ?") parameters.append(kind) if equipment: conditions.append("e.equipment_id = ?") parameters.append(equipment.upper()) fields = ("e.source, e.page, e.kind, e.text, e.pdf_sha256, e.ocr_sha256, e.effective, " "e.equipment_id, e.table_number, e.row_number") if terms: start = f"SELECT {fields} FROM evidence_words JOIN evidence e ON e.id = evidence_words.rowid" conditions.append("evidence_words MATCH ?") parameters.append(terms) order = "ORDER BY bm25(evidence_words), e.id" else: start = f"SELECT {fields} FROM evidence e" order = "ORDER BY e.id" sql = f"{start} WHERE {' AND '.join(conditions)} {order} LIMIT ?" parameters.append(limit) with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as db: return [Evidence(*row) for row in db.execute(sql, parameters)] if __name__ == "__main__": parser = argparse.ArgumentParser(description="Index prose, checked OCR, and versioned table rows.") parser.add_argument("--db", type=Path, default=DEFAULT_DB) commands = parser.add_subparsers(dest="command", required=True) update = commands.add_parser("build") update.add_argument("folder", type=Path) find = commands.add_parser("find") find.add_argument("query", nargs="?", default="") find.add_argument("--kind", choices=KINDS) find.add_argument("--equipment") find.add_argument("--as-of", type=date.fromisoformat, required=True) args = parser.parse_args() try: if args.command == "build": build(args.folder, args.db) else: matches = retrieve(args.db, args.as_of, args.query, args.kind, args.equipment) for item in matches: print(f"{item.citation()}\n {item.text}\n") if not matches: print("No matching evidence.") except (OSError, PdfminerException, PdfReadError, sqlite3.Error, ValueError) as error: raise SystemExit(str(error)) from None