"""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]