356 lines
13 KiB
Python
356 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
文本提取器模块
|
||
支持从PDF、Word、TXT等格式的文档中提取文本内容
|
||
"""
|
||
|
||
import os
|
||
import logging
|
||
from typing import List, Dict, Any, Optional
|
||
from pathlib import Path
|
||
from dataclasses import dataclass
|
||
from datetime import datetime
|
||
|
||
# 导入依赖库
|
||
try:
|
||
import PyPDF2
|
||
import pdfplumber
|
||
PDF_AVAILABLE = True
|
||
except ImportError:
|
||
PDF_AVAILABLE = False
|
||
|
||
try:
|
||
from docx import Document
|
||
DOCX_AVAILABLE = True
|
||
except ImportError:
|
||
DOCX_AVAILABLE = False
|
||
|
||
try:
|
||
import openpyxl
|
||
from openpyxl import load_workbook
|
||
EXCEL_AVAILABLE = True
|
||
except ImportError:
|
||
EXCEL_AVAILABLE = False
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
@dataclass
|
||
class ExtractedDocument:
|
||
"""提取的文档数据"""
|
||
filename: str
|
||
file_type: str
|
||
content: str # 纯文本内容
|
||
metadata: Dict[str, Any] # 文档元数据
|
||
extracted_at: datetime
|
||
file_size: int
|
||
page_count: Optional[int] = None
|
||
|
||
def __post_init__(self):
|
||
# 确保content是字符串
|
||
if not isinstance(self.content, str):
|
||
self.content = str(self.content)
|
||
|
||
class TextExtractor:
|
||
"""文本提取器 - 只做纯文本提取，保留所有原始内容"""
|
||
|
||
def __init__(self, config: Optional[Dict[str, Any]] = None):
|
||
self.config = config or {}
|
||
self.supported_formats = {
|
||
'.pdf': self._extract_pdf,
|
||
'.docx': self._extract_docx,
|
||
'.doc': self._extract_doc,
|
||
'.txt': self._extract_txt,
|
||
'.md': self._extract_txt,
|
||
'.xlsx': self._extract_xlsx,
|
||
'.xls': self._extract_xls,
|
||
'.csv': self._extract_csv
|
||
}
|
||
|
||
def extract(self, file_path: str) -> ExtractedDocument:
|
||
"""提取单个文件的文本内容"""
|
||
path_obj = Path(file_path)
|
||
|
||
if not path_obj.exists():
|
||
raise FileNotFoundError(f"文件不存在: {file_path}")
|
||
|
||
file_ext = path_obj.suffix.lower()
|
||
if file_ext not in self.supported_formats:
|
||
raise ValueError(f"不支持的文件格式: {file_ext}")
|
||
|
||
try:
|
||
# 获取文件信息
|
||
file_size = path_obj.stat().st_size
|
||
|
||
# 提取文本内容
|
||
extractor = self.supported_formats[file_ext]
|
||
content, metadata = extractor(path_obj)
|
||
|
||
return ExtractedDocument(
|
||
filename=path_obj.name,
|
||
file_type=file_ext,
|
||
content=content,
|
||
metadata=metadata,
|
||
extracted_at=datetime.now(),
|
||
file_size=file_size,
|
||
page_count=metadata.get('page_count')
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"提取文件 {file_path} 时出错: {str(e)}")
|
||
raise
|
||
|
||
def extract_batch(self, file_paths: List[str]) -> List[ExtractedDocument]:
|
||
"""批量提取多个文件的文本内容"""
|
||
results = []
|
||
|
||
for file_path in file_paths:
|
||
try:
|
||
result = self.extract(file_path)
|
||
results.append(result)
|
||
logger.info(f"成功提取文件: {file_path}")
|
||
except Exception as e:
|
||
logger.error(f"提取文件 {file_path} 失败: {str(e)}")
|
||
# 创建错误记录
|
||
error_doc = ExtractedDocument(
|
||
filename=Path(file_path).name,
|
||
file_type=Path(file_path).suffix.lower(),
|
||
content=f"提取失败: {str(e)}",
|
||
metadata={"error": str(e)},
|
||
extracted_at=datetime.now(),
|
||
file_size=0
|
||
)
|
||
results.append(error_doc)
|
||
|
||
return results
|
||
|
||
def _extract_pdf(self, file_path: Path) -> tuple[str, Dict[str, Any]]:
|
||
"""提取PDF文件的纯文本内容"""
|
||
if not PDF_AVAILABLE:
|
||
raise ImportError("需要安装 PyPDF2 和 pdfplumber: pip install PyPDF2 pdfplumber")
|
||
|
||
content_parts = []
|
||
metadata = {}
|
||
|
||
try:
|
||
# 使用pdfplumber提取文本（更好的文本提取）
|
||
with pdfplumber.open(file_path) as pdf:
|
||
metadata['page_count'] = len(pdf.pages)
|
||
|
||
for page_num, page in enumerate(pdf.pages, 1):
|
||
page_text = page.extract_text()
|
||
if page_text:
|
||
content_parts.append(f"=== 第 {page_num} 页 ===\n{page_text}\n")
|
||
|
||
# 获取文档元数据
|
||
if pdf.metadata:
|
||
metadata.update({
|
||
'title': pdf.metadata.get('Title', ''),
|
||
'author': pdf.metadata.get('Author', ''),
|
||
'subject': pdf.metadata.get('Subject', ''),
|
||
'creator': pdf.metadata.get('Creator', ''),
|
||
'producer': pdf.metadata.get('Producer', ''),
|
||
'creation_date': pdf.metadata.get('CreationDate', ''),
|
||
'modification_date': pdf.metadata.get('ModDate', '')
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.warning(f"pdfplumber提取失败，尝试使用PyPDF2: {str(e)}")
|
||
|
||
# 备用方案：使用PyPDF2
|
||
with open(file_path, 'rb') as file:
|
||
pdf_reader = PyPDF2.PdfReader(file)
|
||
metadata['page_count'] = len(pdf_reader.pages)
|
||
|
||
for page_num, page in enumerate(pdf_reader.pages, 1):
|
||
page_text = page.extract_text()
|
||
if page_text:
|
||
content_parts.append(f"=== 第 {page_num} 页 ===\n{page_text}\n")
|
||
|
||
# 获取文档元数据
|
||
if pdf_reader.metadata:
|
||
metadata.update({
|
||
'title': pdf_reader.metadata.get('/Title', ''),
|
||
'author': pdf_reader.metadata.get('/Author', ''),
|
||
'subject': pdf_reader.metadata.get('/Subject', ''),
|
||
'creator': pdf_reader.metadata.get('/Creator', ''),
|
||
'producer': pdf_reader.metadata.get('/Producer', ''),
|
||
'creation_date': pdf_reader.metadata.get('/CreationDate', ''),
|
||
'modification_date': pdf_reader.metadata.get('/ModDate', '')
|
||
})
|
||
|
||
content = '\n'.join(content_parts) if content_parts else ""
|
||
return content, metadata
|
||
|
||
def _extract_docx(self, file_path: Path) -> tuple[str, Dict[str, Any]]:
|
||
"""提取DOCX文件的纯文本内容"""
|
||
if not DOCX_AVAILABLE:
|
||
raise ImportError("需要安装 python-docx: pip install python-docx")
|
||
|
||
doc = Document(str(file_path))
|
||
content_parts = []
|
||
metadata = {}
|
||
|
||
# 提取所有段落文本
|
||
for paragraph in doc.paragraphs:
|
||
if paragraph.text.strip():
|
||
content_parts.append(paragraph.text)
|
||
|
||
# 提取表格内容
|
||
for table in doc.tables:
|
||
table_content = []
|
||
for row in table.rows:
|
||
row_content = []
|
||
for cell in row.cells:
|
||
row_content.append(cell.text.strip())
|
||
table_content.append('\t'.join(row_content))
|
||
if table_content:
|
||
content_parts.append('\n=== 表格 ===\n' + '\n'.join(table_content) + '\n')
|
||
|
||
# 获取文档属性
|
||
core_props = doc.core_properties
|
||
metadata.update({
|
||
'title': core_props.title or '',
|
||
'author': core_props.author or '',
|
||
'subject': core_props.subject or '',
|
||
'keywords': core_props.keywords or '',
|
||
'comments': core_props.comments or '',
|
||
'created': str(core_props.created) if core_props.created else '',
|
||
'modified': str(core_props.modified) if core_props.modified else '',
|
||
'last_modified_by': core_props.last_modified_by or '',
|
||
'paragraph_count': len(doc.paragraphs),
|
||
'table_count': len(doc.tables)
|
||
})
|
||
|
||
content = '\n'.join(content_parts)
|
||
return content, metadata
|
||
|
||
def _extract_doc(self, file_path: Path) -> tuple[str, Dict[str, Any]]:
|
||
"""提取DOC文件的纯文本内容"""
|
||
# DOC格式较复杂，建议转换为DOCX或使用专门的库
|
||
logger.warning("DOC格式支持有限，建议转换为DOCX格式")
|
||
|
||
# 尝试读取为文本文件
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
|
||
content = file.read()
|
||
except:
|
||
with open(file_path, 'r', encoding='gbk', errors='ignore') as file:
|
||
content = file.read()
|
||
|
||
metadata = {'format': 'doc', 'encoding_note': '可能存在编码问题'}
|
||
return content, metadata
|
||
|
||
def _extract_txt(self, file_path: Path) -> tuple[str, Dict[str, Any]]:
|
||
"""提取TXT/MD文件的纯文本内容"""
|
||
encodings = ['utf-8', 'gbk', 'gb2312', 'big5', 'utf-16']
|
||
content = ""
|
||
used_encoding = ""
|
||
|
||
for encoding in encodings:
|
||
try:
|
||
with open(file_path, 'r', encoding=encoding) as file:
|
||
content = file.read()
|
||
used_encoding = encoding
|
||
break
|
||
except UnicodeDecodeError:
|
||
continue
|
||
|
||
if not content:
|
||
# 最后尝试忽略错误
|
||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
|
||
content = file.read()
|
||
used_encoding = 'utf-8 (with errors ignored)'
|
||
|
||
metadata = {
|
||
'encoding': used_encoding,
|
||
'line_count': len(content.splitlines()),
|
||
'char_count': len(content)
|
||
}
|
||
|
||
return content, metadata
|
||
|
||
def _extract_xlsx(self, file_path: Path) -> tuple[str, Dict[str, Any]]:
|
||
"""提取XLSX文件的纯文本内容"""
|
||
if not EXCEL_AVAILABLE:
|
||
raise ImportError("需要安装 openpyxl: pip install openpyxl")
|
||
|
||
workbook = load_workbook(file_path, read_only=True)
|
||
content_parts = []
|
||
metadata = {
|
||
'sheet_count': len(workbook.sheetnames),
|
||
'sheet_names': workbook.sheetnames
|
||
}
|
||
|
||
for sheet_name in workbook.sheetnames:
|
||
sheet = workbook[sheet_name]
|
||
content_parts.append(f"\n=== 工作表: {sheet_name} ===\n")
|
||
|
||
for row in sheet.iter_rows(values_only=True):
|
||
row_content = []
|
||
for cell in row:
|
||
if cell is not None:
|
||
row_content.append(str(cell))
|
||
else:
|
||
row_content.append("")
|
||
if any(cell.strip() for cell in row_content): # 跳过空行
|
||
content_parts.append('\t'.join(row_content))
|
||
|
||
content = '\n'.join(content_parts)
|
||
return content, metadata
|
||
|
||
def _extract_xls(self, file_path: Path) -> tuple[str, Dict[str, Any]]:
|
||
"""提取XLS文件的纯文本内容"""
|
||
logger.warning("XLS格式支持有限，建议转换为XLSX格式")
|
||
|
||
# 简单的文本提取
|
||
try:
|
||
with open(file_path, 'rb') as file:
|
||
content = file.read().decode('utf-8', errors='ignore')
|
||
except:
|
||
content = f"无法读取XLS文件: {file_path}"
|
||
|
||
metadata = {'format': 'xls', 'note': '可能存在格式问题'}
|
||
return content, metadata
|
||
|
||
def _extract_csv(self, file_path: Path) -> tuple[str, Dict[str, Any]]:
|
||
"""提取CSV文件的纯文本内容"""
|
||
encodings = ['utf-8', 'gbk', 'gb2312']
|
||
content = ""
|
||
used_encoding = ""
|
||
|
||
for encoding in encodings:
|
||
try:
|
||
with open(file_path, 'r', encoding=encoding) as file:
|
||
content = file.read()
|
||
used_encoding = encoding
|
||
break
|
||
except UnicodeDecodeError:
|
||
continue
|
||
|
||
if not content:
|
||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
|
||
content = file.read()
|
||
used_encoding = 'utf-8 (with errors ignored)'
|
||
|
||
# 计算行数和列数
|
||
lines = content.splitlines()
|
||
row_count = len(lines)
|
||
col_count = len(lines[0].split(',')) if lines else 0
|
||
|
||
metadata = {
|
||
'encoding': used_encoding,
|
||
'row_count': row_count,
|
||
'estimated_col_count': col_count
|
||
}
|
||
|
||
return content, metadata
|
||
|
||
def get_supported_formats(self) -> List[str]:
|
||
"""获取支持的文件格式列表"""
|
||
return list(self.supported_formats.keys())
|
||
|
||
def is_supported(self, file_path: str) -> bool:
|
||
"""检查文件格式是否支持"""
|
||
return Path(file_path).suffix.lower() in self.supported_formats |