"""多格式文件文本提取:按扩展名分发到对应解析器 支持的扩展名: - .txt / .md:UTF-8 解码(errors="replace" 兜底) - .html / .htm:标准库 html.parser 剥离标签提取可见文本 - .pdf:pypdf 逐页 extract_text 拼接;文本层为空(扫描件/图片型)时 自动降级为 OCR(pypdfium2 渲染 + rapidocr-onnxruntime 识别) - .docx:python-docx 段落文本拼接(不含表格/页眉页脚) 未识别扩展名抛 ValueError("不支持的文件类型: {ext}"); 解析异常统一包装为 ValueError("文件解析失败: {detail}"),原异常链式保留。 """ from __future__ import annotations import io from collections.abc import Callable from html.parser import HTMLParser from pathlib import Path from typing import Any import structlog from app.config import settings logger = structlog.get_logger() # 模块级懒加载 OCR 引擎单例(首次调用时初始化,避免无扫描件场景白白下载模型) _ocr_engine: Any | None = None _ocr_unavailable: bool = False # 标记 OCR 依赖不可用,后续直接跳过避免重复尝试 class _VisibleTextExtractor(HTMLParser): """HTMLParser 子类:累积可见文本,跳过 script/style 内容""" def __init__(self) -> None: super().__init__(convert_charrefs=True) self._parts: list[str] = [] self._skip_depth = 0 # 在 script/style 标签内时 > 0 def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: if tag.lower() in {"script", "style"}: self._skip_depth += 1 def handle_endtag(self, tag: str) -> None: if tag.lower() in {"script", "style"} and self._skip_depth > 0: self._skip_depth -= 1 def handle_data(self, data: str) -> None: if self._skip_depth == 0: self._parts.append(data) def get_text(self) -> str: # 块级标签间用空格连接,再折叠多余空白 text = " ".join(self._parts) return " ".join(text.split()) def _decode_html(content: bytes) -> str: """HTML 内容解码:优先 utf-8(带 BOM),失败回退 latin-1""" try: return content.decode("utf-8-sig") except UnicodeDecodeError: return content.decode("latin-1", errors="replace") def _parse_text(content: bytes) -> str: """UTF-8 解码(errors=replace 兜底),保留原字符""" return content.decode("utf-8", errors="replace") def _parse_html(content: bytes) -> str: """HTML 剥离标签,保留可见文本""" parser = _VisibleTextExtractor() parser.feed(_decode_html(content)) parser.close() return parser.get_text() def _parse_pdf(content: bytes) -> str: """PDF 解析:优先 pypdf extract_text;文本层为空(扫描件)时降级 OCR OCR 流程:pypdfium2 渲染每页为 PIL Image → rapidocr-onnxruntime 识别 → 拼接每页识别出的文本。受 settings.pdf_ocr_* 控制:开关、最大页数、DPI。 OCR 依赖未安装或运行异常时降级返回空字符串(由上游 upload 端点拒绝入库)。 """ from pypdf import PdfReader reader = PdfReader(io.BytesIO(content)) parts: list[str] = [] for page in reader.pages: text = page.extract_text() or "" if text: parts.append(text) text_layer = "\n".join(parts).strip() # 文本层非空:直接返回 if text_layer: return text_layer # 文本层为空 → 尝试 OCR 降级 if not settings.pdf_ocr_enabled: return "" ocr_text = _ocr_pdf( content, max_pages=settings.pdf_ocr_max_pages, dpi=settings.pdf_ocr_dpi ) return ocr_text def _ocr_pdf(content: bytes, max_pages: int, dpi: int) -> str: """对扫描件 PDF 跑 OCR:渲染每页 → 识别 → 拼接 返回空字符串的场景:依赖未安装 / 渲染或识别异常 / 无识别结果。 任何异常仅告警不抛出,由上游按"无法提取文本"处理。 """ global _ocr_engine, _ocr_unavailable if _ocr_unavailable: return "" # 1. 懒加载 OCR 引擎 if _ocr_engine is None: try: from rapidocr_onnxruntime import RapidOCR _ocr_engine = RapidOCR() logger.info("PDF OCR 引擎已初始化", dpi=dpi) except Exception: _ocr_unavailable = True logger.warning("OCR 依赖不可用,扫描件 PDF 将无法提取文本", exc_info=True) return "" # 2. 渲染并识别 try: import pypdfium2 as pdfium scale = max(1.0, dpi / 72.0) pdf = pdfium.PdfDocument(io.BytesIO(content)) total = min(len(pdf), max(1, max_pages)) page_texts: list[str] = [] for i in range(total): page = pdf[i] pil_image = page.render(scale=scale).to_pil() result, _ = _ocr_engine(pil_image) if result: # result: [[box, text, score], ...],按行拼接 lines = [ item[1] for item in result if item and len(item) >= 2 and item[1] ] if lines: page_texts.append("\n".join(lines)) pdf.close() return "\n".join(page_texts).strip() except Exception as exc: logger.warning("PDF OCR 失败,降级返回空文本", error=str(exc), exc_info=True) return "" def _parse_docx(content: bytes) -> str: """DOCX 段落文本拼接(不含表格/页眉页脚)""" from docx import Document # type: ignore[import-untyped] document = Document(io.BytesIO(content)) parts = [p.text for p in document.paragraphs if p.text and p.text.strip()] return "\n".join(parts).strip() # 扩展名 → 解析函数映射(启动时构建,避免每次请求重复构造) _PARSERS: dict[str, Callable[[bytes], str]] = { ".txt": _parse_text, ".md": _parse_text, ".html": _parse_html, ".htm": _parse_html, ".pdf": _parse_pdf, ".docx": _parse_docx, } def parse_file(filename: str, content: bytes) -> str: """按扩展名分发解析器提取文本 Args: filename: 文件名(用于判定扩展名) content: 文件二进制内容 Returns: 提取的纯文本 Raises: ValueError: 未识别扩展名 → "不支持的文件类型: {ext}" 解析失败 → "文件解析失败: {detail}"(保留原异常链) """ ext = Path(filename).suffix.lower() parser = _PARSERS.get(ext) if parser is None: raise ValueError(f"不支持的文件类型: {ext or '(无扩展名)'}") try: return parser(content) except ValueError: raise except Exception as exc: raise ValueError(f"文件解析失败: {exc}") from exc def supported_extensions() -> set[str]: """返回当前支持的扩展名集合(含点号,全小写)""" return set(_PARSERS.keys())