Réduit les erreurs ruff de 11 à 5 (toutes pré-existantes non auto-fixables) : - F401 imports inutilisés dans auth/mfa.py - I001 blocs d'imports non triés dans auth/router.py + main.py + pdf_reader.py Les 5 restantes sont dans bookslm/export/watcher (code pré-existant, hors scope). Vérifié : pytest full suite reste 492 passed, 5 skipped, 0 failed.
131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
"""PDF text extraction for ObsiGate — pypdf backend, pymupdf optional fallback."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
PDF_READER: str = "pypdf"
|
|
PdfReader = None # type: ignore
|
|
try:
|
|
import fitz # pymupdf
|
|
PDF_READER = "pymupdf"
|
|
logger.info("PDF reader: pymupdf (high performance)")
|
|
except ImportError:
|
|
try:
|
|
from pypdf import PdfReader # type: ignore
|
|
logger.info("PDF reader: pypdf (pure Python)")
|
|
except ImportError:
|
|
logger.warning("No PDF reader available — install pypdf or pymupdf")
|
|
|
|
|
|
def extract_pdf_text(file_path: Path, max_chars: int = 100000) -> str:
|
|
"""Extract text from a PDF file. Returns empty string on failure."""
|
|
if PdfReader is None and PDF_READER == "pypdf":
|
|
return ""
|
|
try:
|
|
if PDF_READER == "pymupdf":
|
|
return _extract_pymupdf(file_path, max_chars)
|
|
else:
|
|
return _extract_pypdf(file_path, max_chars)
|
|
except Exception as e:
|
|
logger.warning("Failed to extract PDF text from %s: %s", file_path, e)
|
|
return ""
|
|
|
|
|
|
def extract_pdf_metadata(file_path: Path) -> dict:
|
|
"""Extract metadata from a PDF file."""
|
|
info = {"pages": 0, "title": "", "author": ""}
|
|
try:
|
|
if PDF_READER == "pymupdf":
|
|
doc = fitz.open(str(file_path))
|
|
info["pages"] = doc.page_count
|
|
meta = doc.metadata or {}
|
|
info["title"] = meta.get("title", "")
|
|
info["author"] = meta.get("author", "")
|
|
doc.close()
|
|
else:
|
|
reader = PdfReader(str(file_path))
|
|
info["pages"] = len(reader.pages)
|
|
meta = reader.metadata or {}
|
|
if meta:
|
|
info["title"] = str(meta.get("/Title", ""))
|
|
info["author"] = str(meta.get("/Author", ""))
|
|
except Exception as e:
|
|
logger.warning("Failed to extract PDF metadata from %s: %s", file_path, e)
|
|
return info
|
|
|
|
|
|
def extract_pdf_toc(file_path: Path) -> list[dict]:
|
|
"""Extract table of contents (bookmarks/outline) from a PDF file.
|
|
|
|
Returns a list of {title, page, level} dicts, or empty list on failure.
|
|
"""
|
|
toc: list[dict] = []
|
|
try:
|
|
if PDF_READER == "pymupdf":
|
|
doc = fitz.open(str(file_path))
|
|
raw = doc.get_toc(simple=False)
|
|
for item in raw:
|
|
toc.append({
|
|
"title": str(item[1]),
|
|
"page": item[2],
|
|
"level": item[0],
|
|
})
|
|
doc.close()
|
|
elif PdfReader is not None:
|
|
reader = PdfReader(str(file_path))
|
|
outline = reader.outline
|
|
if outline:
|
|
def _flatten(items, level=1):
|
|
for item in items:
|
|
if isinstance(item, list):
|
|
_flatten(item, level + 1)
|
|
elif hasattr(item, 'title') and hasattr(item, 'page'):
|
|
page_num = reader.get_page_number(item.page) + 1 if hasattr(item, 'page') and item.page else 1
|
|
toc.append({
|
|
"title": str(item.title),
|
|
"page": page_num,
|
|
"level": level,
|
|
})
|
|
_flatten(outline)
|
|
except Exception as e:
|
|
logger.warning("Failed to extract PDF TOC from %s: %s", file_path, e)
|
|
return toc
|
|
|
|
|
|
def _extract_pymupdf(file_path: Path, max_chars: int) -> str:
|
|
doc = fitz.open(str(file_path))
|
|
parts = []
|
|
total = 0
|
|
for page in doc:
|
|
text = page.get_text()
|
|
if text:
|
|
if total + len(text) > max_chars:
|
|
remaining = max_chars - total
|
|
if remaining > 0:
|
|
parts.append(text[:remaining])
|
|
break
|
|
parts.append(text)
|
|
total += len(text)
|
|
doc.close()
|
|
return "\f".join(parts)
|
|
|
|
|
|
def _extract_pypdf(file_path: Path, max_chars: int) -> str:
|
|
reader = PdfReader(str(file_path))
|
|
parts = []
|
|
total = 0
|
|
for page in reader.pages:
|
|
text = page.extract_text() or ""
|
|
if text:
|
|
if total + len(text) > max_chars:
|
|
remaining = max_chars - total
|
|
if remaining > 0:
|
|
parts.append(text[:remaining])
|
|
break
|
|
parts.append(text)
|
|
total += len(text)
|
|
return "\f".join(parts)
|