113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
"""PDF text extraction service using PyMuPDF and pdfplumber."""
|
|
import fitz # PyMuPDF
|
|
import pdfplumber
|
|
import logging
|
|
from typing import Tuple, Optional
|
|
from io import BytesIO
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PDFExtractionError(Exception):
|
|
"""Custom exception for PDF extraction errors."""
|
|
pass
|
|
|
|
|
|
def extract_text_with_pymupdf(pdf_bytes: bytes) -> Tuple[str, int]:
|
|
"""
|
|
Extract text from PDF using PyMuPDF (fitz).
|
|
|
|
Args:
|
|
pdf_bytes: PDF file content as bytes
|
|
|
|
Returns:
|
|
Tuple of (extracted_text, page_count)
|
|
|
|
Raises:
|
|
PDFExtractionError: If extraction fails
|
|
"""
|
|
try:
|
|
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
|
|
|
|
if doc.is_encrypted:
|
|
raise PDFExtractionError("PDF is encrypted and cannot be read")
|
|
|
|
page_count = len(doc)
|
|
text_parts = []
|
|
|
|
for page in doc:
|
|
text_parts.append(page.get_text())
|
|
|
|
doc.close()
|
|
full_text = "\n".join(text_parts)
|
|
|
|
logger.info(f"Extracted {len(full_text)} characters using PyMuPDF from {page_count} pages")
|
|
return full_text, page_count
|
|
|
|
except Exception as e:
|
|
logger.warning(f"PyMuPDF extraction failed: {str(e)}")
|
|
raise PDFExtractionError(f"PyMuPDF extraction failed: {str(e)}")
|
|
|
|
|
|
def extract_text_with_pdfplumber(pdf_bytes: bytes) -> Tuple[str, int]:
|
|
"""
|
|
Extract text from PDF using pdfplumber (fallback method).
|
|
|
|
Args:
|
|
pdf_bytes: PDF file content as bytes
|
|
|
|
Returns:
|
|
Tuple of (extracted_text, page_count)
|
|
|
|
Raises:
|
|
PDFExtractionError: If extraction fails
|
|
"""
|
|
try:
|
|
with pdfplumber.open(BytesIO(pdf_bytes)) as pdf:
|
|
page_count = len(pdf.pages)
|
|
text_parts = []
|
|
|
|
for page in pdf.pages:
|
|
page_text = page.extract_text()
|
|
if page_text:
|
|
text_parts.append(page_text)
|
|
|
|
full_text = "\n".join(text_parts)
|
|
|
|
logger.info(f"Extracted {len(full_text)} characters using pdfplumber from {page_count} pages")
|
|
return full_text, page_count
|
|
|
|
except Exception as e:
|
|
logger.error(f"pdfplumber extraction failed: {str(e)}")
|
|
raise PDFExtractionError(f"pdfplumber extraction failed: {str(e)}")
|
|
|
|
|
|
def extract_text_from_pdf(pdf_bytes: bytes) -> Tuple[str, int, str]:
|
|
"""
|
|
Extract text from PDF using available methods.
|
|
Tries PyMuPDF first, falls back to pdfplumber.
|
|
|
|
Args:
|
|
pdf_bytes: PDF file content as bytes
|
|
|
|
Returns:
|
|
Tuple of (extracted_text, page_count, method_used)
|
|
|
|
Raises:
|
|
PDFExtractionError: If all extraction methods fail
|
|
"""
|
|
# Try PyMuPDF first
|
|
try:
|
|
text, pages = extract_text_with_pymupdf(pdf_bytes)
|
|
return text, pages, "text"
|
|
except PDFExtractionError as e:
|
|
logger.warning(f"PyMuPDF failed, trying pdfplumber: {str(e)}")
|
|
|
|
# Fallback to pdfplumber
|
|
try:
|
|
text, pages = extract_text_with_pdfplumber(pdf_bytes)
|
|
return text, pages, "text"
|
|
except PDFExtractionError as e:
|
|
logger.error(f"All extraction methods failed: {str(e)}")
|
|
raise PDFExtractionError("Failed to extract text from PDF using all available methods")
|