111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
"""
|
|
Excel / PDF 抽出纯文本。
|
|
|
|
本文件职责：把允许的附件字节变成可见原文，供千问整理、再走文字询价。
|
|
属于 policy 包。禁止在这里猜港口/箱型/缺项；禁止写六态。
|
|
识别槽与 H5 都只调用本模块的纯函数。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import logging
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_ALLOWED_EXT = {".xls", ".xlsx", ".pdf"}
|
|
_REJECT_HINT = "only_xlsx_xls_pdf"
|
|
|
|
|
|
def attachment_ext(filename: str) -> str:
|
|
"""小写后缀，含点。无后缀返回空。"""
|
|
name = (filename or "").strip().rsplit("/", 1)[-1].rsplit("\\", 1)[-1]
|
|
if "." not in name:
|
|
return ""
|
|
return "." + name.rsplit(".", 1)[-1].lower()
|
|
|
|
|
|
def allowed_attachment(filename: str) -> bool:
|
|
"""H5 只收 xls / xlsx / pdf。"""
|
|
return attachment_ext(filename) in _ALLOWED_EXT
|
|
|
|
|
|
def extract_attachment_text(filename: str, raw: bytes) -> dict[str, Any]:
|
|
"""
|
|
抽出单元格或 PDF 可见字。
|
|
|
|
返回 ok/text/error/empty。失败时 text 为空，由上层走「没看清 / 暂时没认出」。
|
|
"""
|
|
ext = attachment_ext(filename)
|
|
if ext not in _ALLOWED_EXT:
|
|
return {"ok": False, "text": "", "empty": True, "error": _REJECT_HINT}
|
|
blob = raw or b""
|
|
if not blob:
|
|
return {"ok": False, "text": "", "empty": True, "error": "empty_bytes"}
|
|
try:
|
|
if ext == ".pdf":
|
|
text = _pdf_text(blob)
|
|
elif ext == ".xlsx":
|
|
text = _xlsx_text(blob)
|
|
else:
|
|
text = _xls_text(blob)
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.exception("抽附件文本失败 name=%s", filename)
|
|
return {"ok": False, "text": "", "empty": True, "error": f"extract_fail:{type(exc).__name__}"}
|
|
body = (text or "").strip()
|
|
return {"ok": True, "text": body, "empty": not bool(body), "error": ""}
|
|
|
|
|
|
def _xlsx_text(raw: bytes) -> str:
|
|
from openpyxl import load_workbook
|
|
|
|
wb = load_workbook(io.BytesIO(raw), read_only=True, data_only=True)
|
|
parts: list[str] = []
|
|
try:
|
|
for sheet in wb.worksheets:
|
|
title = str(sheet.title or "").strip()
|
|
if title:
|
|
parts.append(f"# {title}")
|
|
for row in sheet.iter_rows(values_only=True):
|
|
cells = [str(c).strip() for c in row if c is not None and str(c).strip()]
|
|
if cells:
|
|
parts.append("\t".join(cells))
|
|
finally:
|
|
wb.close()
|
|
return "\n".join(parts)
|
|
|
|
|
|
def _xls_text(raw: bytes) -> str:
|
|
try:
|
|
import xlrd
|
|
except ImportError:
|
|
return ""
|
|
book = xlrd.open_workbook(file_contents=raw)
|
|
parts: list[str] = []
|
|
for sheet in book.sheets():
|
|
title = str(sheet.name or "").strip()
|
|
if title:
|
|
parts.append(f"# {title}")
|
|
for r in range(sheet.nrows):
|
|
cells = [
|
|
str(sheet.cell_value(r, c)).strip()
|
|
for c in range(sheet.ncols)
|
|
if str(sheet.cell_value(r, c)).strip()
|
|
]
|
|
if cells:
|
|
parts.append("\t".join(cells))
|
|
return "\n".join(parts)
|
|
|
|
|
|
def _pdf_text(raw: bytes) -> str:
|
|
from pypdf import PdfReader
|
|
|
|
reader = PdfReader(io.BytesIO(raw))
|
|
parts: list[str] = []
|
|
for page in reader.pages:
|
|
piece = page.extract_text() or ""
|
|
if piece.strip():
|
|
parts.append(piece.strip())
|
|
return "\n".join(parts)
|