58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
"""
|
|
Tabular-family extractor: csv.
|
|
|
|
Renders to a markdown table (capped preview) with row/column metadata. Uses only
|
|
the stdlib so it always works.
|
|
"""
|
|
import csv
|
|
from pathlib import Path
|
|
|
|
from .base import ExtractionResult
|
|
|
|
EXTRACTOR_ID = "tabular/csv@1"
|
|
MAX_PREVIEW_ROWS = 200
|
|
|
|
|
|
def _md_table(rows: list) -> str:
|
|
if not rows:
|
|
return ""
|
|
width = max(len(r) for r in rows)
|
|
rows = [r + [""] * (width - len(r)) for r in rows]
|
|
header, body = rows[0], rows[1:]
|
|
|
|
def esc(c: str) -> str:
|
|
return str(c).replace("|", "\\|").replace("\n", " ")
|
|
|
|
lines = ["| " + " | ".join(esc(c) for c in header) + " |",
|
|
"| " + " | ".join(["---"] * width) + " |"]
|
|
for r in body:
|
|
lines.append("| " + " | ".join(esc(c) for c in r) + " |")
|
|
return "\n".join(lines) + "\n"
|
|
|
|
|
|
def extract(source: Path, options: dict) -> ExtractionResult:
|
|
warnings = []
|
|
with open(source, newline="", encoding="utf-8", errors="replace") as f:
|
|
sample = f.read(8192)
|
|
f.seek(0)
|
|
try:
|
|
dialect = csv.Sniffer().sniff(sample) if sample.strip() else csv.excel
|
|
except csv.Error:
|
|
dialect = csv.excel
|
|
rows = list(csv.reader(f, dialect))
|
|
|
|
total = len(rows)
|
|
cols = max((len(r) for r in rows), default=0)
|
|
preview = rows[: MAX_PREVIEW_ROWS + 1] # +1 for header
|
|
if total > MAX_PREVIEW_ROWS + 1:
|
|
warnings.append(f"table has {total} rows; previewed first {MAX_PREVIEW_ROWS}")
|
|
|
|
content = _md_table(preview)
|
|
metadata = {
|
|
"rows": max(total - 1, 0),
|
|
"columns": cols,
|
|
"header": rows[0] if rows else [],
|
|
}
|
|
return ExtractionResult(content=content, metadata=metadata,
|
|
warnings=warnings, extractor=EXTRACTOR_ID)
|