Blame for finbot/extract.py

"""Turn an uploaded file into text and rows.

Deliberately format-agnostic and bank-agnostic: this module only gets the
content *out* of the file. Deciding which column is a date and which is an
amount is :mod:`finbot.ingest`'s job, and is done by the model so that any
bank's export works without a hand-written layout per institution.

Nothing here calls the network, so all of it is unit-testable offline.
"""

from __future__ import annotations

import csv
import io
import logging
from dataclasses import dataclass, field

log = logging.getLogger(__name__)

#: Guard rails. A personal statement is a few hundred rows; anything far past
#: that is either a mistake or would blow the model's context and the user's
#: token budget, so it is truncated with a visible note rather than silently.
MAX_ROWS = 1500
MAX_TEXT_CHARS = 120_000
MAX_PDF_PAGES = 40


@dataclass
class ExtractedDocument:
    kind: str
    text: str = ""
    rows: list[list[str]] = field(default_factory=list)
    pages: int = 0
    truncated: bool = False
    note: str = ""

    @property
    def is_empty(self) -> bool:
        return not self.text.strip() and not self.rows


def _decode(data: bytes) -> str:
    for encoding in ("utf-8-sig", "utf-8", "cp1252", "latin-1"):
        try:
            return data.decode(encoding)
        except UnicodeDecodeError:
            continue
    return data.decode("utf-8", errors="replace")


def _clean_row(row: list) -> list[str]:
    return [("" if c is None else str(c)).strip() for c in row]


def _is_blank(row: list[str]) -> bool:
    return not any(cell for cell in row)


def extract_csv(data: bytes) -> ExtractedDocument:
    text = _decode(data)
    sample = text[:8192]

    try:
        dialect = csv.Sniffer().sniff(sample, delimiters=",;\t|")
    except csv.Error:
        # Sniffing fails on single-column or irregular exports; comma is the
        # safe default and the model tolerates a wrong split better than a crash.
        dialect = csv.excel

    rows: list[list[str]] = []
    truncated = False
    for row in csv.reader(io.StringIO(text), dialect):
        cleaned = _clean_row(row)
        if _is_blank(cleaned):
            continue
        rows.append(cleaned)
        if len(rows) >= MAX_ROWS:
            truncated = True
            break

    return ExtractedDocument(kind="csv", rows=rows, truncated=truncated)


def extract_xlsx(data: bytes) -> ExtractedDocument:
    from openpyxl import load_workbook

    workbook = load_workbook(io.BytesIO(data), read_only=True, data_only=True)
    rows: list[list[str]] = []
    truncated = False

    try:
        for sheet in workbook.worksheets:
            if truncated:
                break
            if len(workbook.worksheets) > 1:
                rows.append([f"# sheet: {sheet.title}"])
            for row in sheet.iter_rows(values_only=True):
                cleaned = _clean_row(list(row))
                if _is_blank(cleaned):
                    continue
                rows.append(cleaned)
                if len(rows) >= MAX_ROWS:
                    truncated = True
                    break
    finally:
        workbook.close()

    return ExtractedDocument(kind="xlsx", rows=rows, truncated=truncated)


def extract_pdf(data: bytes) -> ExtractedDocument:
    """Pull tables (preferred) and text out of a PDF statement.

    Scanned statements contain no text layer. That is detected and reported
    rather than silently returning nothing, because "0 transactions found" and
    "this file is a photograph of a statement" need very different user advice.
    """
    import pdfplumber

    rows: list[list[str]] = []
    text_parts: list[str] = []
    truncated = False
    page_count = 0

    with pdfplumber.open(io.BytesIO(data)) as pdf:
        page_count = len(pdf.pages)
        for index, page in enumerate(pdf.pages):
            if index >= MAX_PDF_PAGES:
                truncated = True
                break

            try:
                for table in page.extract_tables() or []:
                    for row in table:
                        cleaned = _clean_row(list(row))
                        if _is_blank(cleaned):
                            continue
                        rows.append(cleaned)
                        if len(rows) >= MAX_ROWS:
                            truncated = True
                            break
                    if truncated:
                        break
            except Exception as exc:
                log.warning("Table extraction failed on page %d: %s", index + 1, exc)

            try:
                if page_text := page.extract_text():
                    text_parts.append(page_text)
            except Exception as exc:
                log.warning("Text extraction failed on page %d: %s", index + 1, exc)

            if truncated:
                break

    text = "\n".join(text_parts)[:MAX_TEXT_CHARS]
    doc = ExtractedDocument(
        kind="pdf", text=text, rows=rows, pages=page_count, truncated=truncated
    )

    if doc.is_empty:
        doc.note = (
            "This PDF has no selectable text, so it is almost certainly a scan "
            "or photo of a statement. Send it as a photo instead and I'll read "
            "it with the vision model, or export a CSV/XLSX from your bank."
        )
    return doc


def extract(data: bytes, filename: str = "", mime_type: str = "") -> ExtractedDocument:
    """Dispatch on filename extension, falling back to MIME type."""
    name = (filename or "").lower()
    mime = (mime_type or "").lower()

    if name.endswith(".csv") or name.endswith(".tsv") or "csv" in mime:
        return extract_csv(data)
    if name.endswith((".xlsx", ".xlsm", ".xltx")) or "spreadsheet" in mime or "excel" in mime:
        return extract_xlsx(data)
    if name.endswith(".pdf") or "pdf" in mime:
        return extract_pdf(data)
    if name.endswith((".txt", ".md")) or mime.startswith("text/"):
        return ExtractedDocument(kind="text", text=_decode(data)[:MAX_TEXT_CHARS])

    # .xls (old binary format) is not readable by openpyxl -- say so precisely.
    if name.endswith(".xls"):
        return ExtractedDocument(
            kind="unsupported",
            note=(
                "Old .xls files aren't supported. Open it and re-save as .xlsx "
                "or export CSV, then send it again."
            ),
        )

    return ExtractedDocument(
        kind="unsupported",
        note=f"I can't read {filename or 'that file'}. Send a CSV, XLSX, PDF or a photo.",
    )


def rows_to_text(rows: list[list[str]], limit: int = MAX_ROWS) -> str:
    """Render rows as pipe-delimited lines for a model prompt."""
    lines = [" | ".join(cell for cell in row) for row in rows[:limit]]
    return "\n".join(lines)[:MAX_TEXT_CHARS]