cryptorugmunch
47ba268131
Adds missing standard artifacts: - README.md (if missing) - AGENTS.md (AI agent contract) - PLAN.md (current sprint) - STATUS.md (where we are) - DEVELOPMENT.md (dev workflow) - DEPLOYMENT.md (deploy procedure) - TESTING.md (test strategy) - DECISIONS.md (ADR index + templates) - .github/CODEOWNERS - .github/workflows/ci.yml Preserves all existing artifacts. Refs: RugMunchMedia/fleet-template
140 lines
4.4 KiB
Python
140 lines
4.4 KiB
Python
"""Pry — PDF Table Extraction using multiple methods.
|
|
Extracts structured tables from PDF documents (financial reports, invoices, etc.)
|
|
Uses pdfplumber, camelot, and pdfminer as fallback methods."""
|
|
|
|
import io
|
|
import logging
|
|
import os
|
|
from contextlib import suppress
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Try different PDF libraries
|
|
_pdfplumber: bool = False
|
|
_camelot: bool = False
|
|
_pdfminer: bool = False
|
|
|
|
try:
|
|
import pdfplumber # noqa: F401
|
|
|
|
_pdfplumber = True
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
import camelot # noqa: F401
|
|
|
|
_camelot = True
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
from pdfminer.high_level import extract_text # noqa: F401
|
|
|
|
_pdfminer = True
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
class PDFTableExtractor:
|
|
"""Extract tables from PDF documents using multiple methods."""
|
|
|
|
def __init__(self, prefer_method: str = "pdfplumber"):
|
|
self.prefer_method = prefer_method
|
|
|
|
def extract(self, pdf_bytes: bytes, method: str = "") -> dict[str, Any]:
|
|
"""Extract tables from a PDF document.
|
|
|
|
Returns: {tables: [...], text: "...", page_count: N, method_used: "..."}
|
|
"""
|
|
method = method or self.prefer_method
|
|
|
|
# Try preferred method first
|
|
if method == "pdfplumber" and _pdfplumber:
|
|
return self._extract_pdfplumber(pdf_bytes)
|
|
if method == "camelot" and _camelot:
|
|
return self._extract_camelot(pdf_bytes)
|
|
if method == "pdfminer" and _pdfminer:
|
|
return self._extract_pdfminer(pdf_bytes)
|
|
|
|
# Fallback chain
|
|
for fallback in ["pdfplumber", "camelot", "pdfminer"]:
|
|
if fallback == "pdfplumber" and _pdfplumber:
|
|
return self._extract_pdfplumber(pdf_bytes)
|
|
if fallback == "camelot" and _camelot:
|
|
return self._extract_camelot(pdf_bytes)
|
|
if fallback == "pdfminer" and _pdfminer:
|
|
return self._extract_pdfminer(pdf_bytes)
|
|
|
|
return {
|
|
"error": "No PDF library available. Install pdfplumber, camelot, or pdfminer.six.",
|
|
"tables": [],
|
|
"text": "",
|
|
}
|
|
|
|
def _extract_pdfplumber(self, pdf_bytes: bytes) -> dict[str, Any]:
|
|
import pdfplumber
|
|
|
|
tables: list[dict[str, Any]] = []
|
|
text = ""
|
|
page_count = 0
|
|
with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf:
|
|
for page in pdf.pages:
|
|
page_count += 1
|
|
text += (page.extract_text() or "") + "\n"
|
|
for table in page.extract_tables():
|
|
if table:
|
|
tables.append(
|
|
{
|
|
"page": page.page_number,
|
|
"rows": table,
|
|
"row_count": len(table),
|
|
"col_count": len(table[0]) if table else 0,
|
|
}
|
|
)
|
|
return {
|
|
"tables": tables,
|
|
"text": text,
|
|
"page_count": page_count,
|
|
"table_count": len(tables),
|
|
"method_used": "pdfplumber",
|
|
}
|
|
|
|
def _extract_camelot(self, pdf_bytes: bytes) -> dict[str, Any]:
|
|
import camelot
|
|
|
|
tables: list[dict[str, Any]] = []
|
|
tmp_path = "/tmp/_pry_pdf.pdf"
|
|
try:
|
|
with open(tmp_path, "wb") as f:
|
|
f.write(pdf_bytes)
|
|
camelot_tables = camelot.read_pdf(tmp_path, pages="all")
|
|
for i, table in enumerate(camelot_tables):
|
|
tables.append(
|
|
{
|
|
"page": i + 1,
|
|
"rows": table.df.values.tolist(),
|
|
"row_count": len(table.df),
|
|
"col_count": len(table.df.columns),
|
|
"accuracy": table.accuracy,
|
|
"whitespace": table.whitespace,
|
|
}
|
|
)
|
|
finally:
|
|
with suppress(OSError):
|
|
os.unlink(tmp_path)
|
|
return {
|
|
"tables": tables,
|
|
"table_count": len(tables),
|
|
"method_used": "camelot",
|
|
}
|
|
|
|
def _extract_pdfminer(self, pdf_bytes: bytes) -> dict[str, Any]:
|
|
from pdfminer.high_level import extract_text
|
|
|
|
text = extract_text(io.BytesIO(pdf_bytes))
|
|
return {"tables": [], "text": text, "method_used": "pdfminer"}
|
|
|
|
def is_available(self) -> bool:
|
|
return any([_pdfplumber, _camelot, _pdfminer])
|