816 lines
20 KiB
Python
816 lines
20 KiB
Python
import io
|
|
import os
|
|
import re
|
|
import csv
|
|
import tempfile
|
|
import subprocess
|
|
import logging
|
|
|
|
from typing import List
|
|
|
|
from bs4 import BeautifulSoup
|
|
|
|
from pypdf import PdfReader
|
|
from docx import Document
|
|
|
|
import openpyxl
|
|
import pdfplumber
|
|
|
|
from PIL import Image
|
|
import pytesseract
|
|
|
|
from striprtf.striprtf import rtf_to_text
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DocumentProcessor:
|
|
|
|
# =========================================================
|
|
# HELPERS
|
|
# =========================================================
|
|
|
|
@staticmethod
|
|
def _cleanup_text(text: str) -> str:
|
|
|
|
if not text:
|
|
return ""
|
|
|
|
text = text.replace("\xa0", " ")
|
|
text = text.replace("\t", " ")
|
|
|
|
text = re.sub(r"[ \t]+", " ", text)
|
|
text = re.sub(r"\n{3,}", "\n\n", text)
|
|
|
|
return text.strip()
|
|
|
|
@staticmethod
|
|
def _cell_to_text(v) -> str:
|
|
|
|
if v is None:
|
|
return ""
|
|
|
|
s = str(v).strip()
|
|
|
|
return s if s else ""
|
|
|
|
# =========================================================
|
|
# HTML
|
|
# =========================================================
|
|
|
|
@staticmethod
|
|
def normalize_email_html(html_content: str) -> str:
|
|
"""
|
|
Очищает HTML письма и нормализует таблицы
|
|
"""
|
|
|
|
soup = BeautifulSoup(
|
|
html_content,
|
|
"html.parser"
|
|
)
|
|
|
|
for tag in soup(["script", "style"]):
|
|
tag.decompose()
|
|
|
|
for table in soup.find_all("table"):
|
|
|
|
table["style"] = (
|
|
"border-collapse: collapse; width: 100%;"
|
|
)
|
|
|
|
for cell in table.find_all(["td", "th"]):
|
|
|
|
cell["style"] = (
|
|
"border:1px solid #ccc;padding:6px;"
|
|
)
|
|
|
|
return str(soup)
|
|
|
|
# =========================================================
|
|
# MAIN
|
|
# =========================================================
|
|
|
|
def extract_text(
|
|
self,
|
|
content: bytes,
|
|
filename: str
|
|
) -> str:
|
|
|
|
ext = filename.lower().split(".")[-1]
|
|
|
|
try:
|
|
|
|
# =================================================
|
|
# PDF
|
|
# =================================================
|
|
|
|
if ext == "pdf":
|
|
return self._extract_pdf(content)
|
|
|
|
# =================================================
|
|
# WORD
|
|
# =================================================
|
|
|
|
elif ext in ["docx", "docm"]:
|
|
return self._extract_docx(content)
|
|
|
|
elif ext == "doc":
|
|
return self._extract_doc(content)
|
|
|
|
elif ext == "rtf":
|
|
return self._extract_rtf(content)
|
|
|
|
# =================================================
|
|
# EXCEL
|
|
# =================================================
|
|
|
|
elif ext in ["xlsx", "xlsm", "xlsxm"]:
|
|
return self._extract_excel_openpyxl(content)
|
|
|
|
elif ext == "xls":
|
|
return self._extract_excel_xls(content)
|
|
|
|
elif ext == "csv":
|
|
return self._extract_csv(content)
|
|
|
|
# =================================================
|
|
# IMAGES / OCR
|
|
# =================================================
|
|
|
|
elif ext in [
|
|
"png",
|
|
"jpg",
|
|
"jpeg",
|
|
"gif",
|
|
"bmp",
|
|
"webp"
|
|
]:
|
|
return self._extract_image(content)
|
|
|
|
elif ext in ["tif", "tiff"]:
|
|
return self._extract_tiff(content)
|
|
|
|
# =================================================
|
|
# TEXT
|
|
# =================================================
|
|
|
|
elif ext == "txt":
|
|
|
|
try:
|
|
return content.decode("utf-8")
|
|
|
|
except Exception:
|
|
return content.decode(
|
|
"cp1251",
|
|
errors="ignore"
|
|
)
|
|
|
|
# =================================================
|
|
# UNKNOWN
|
|
# =================================================
|
|
|
|
logger.warning(
|
|
f"Unsupported attachment format: {ext}"
|
|
)
|
|
|
|
return ""
|
|
|
|
except Exception:
|
|
|
|
logger.exception(
|
|
f"Attachment processing failed: {filename}"
|
|
)
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# PDF
|
|
# =========================================================
|
|
|
|
def _extract_pdf(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
parts = []
|
|
|
|
with pdfplumber.open(
|
|
io.BytesIO(content)
|
|
) as pdf:
|
|
|
|
for page_num, page in enumerate(
|
|
pdf.pages,
|
|
start=1
|
|
):
|
|
|
|
parts.append(
|
|
f"=== PAGE {page_num} ==="
|
|
)
|
|
|
|
# =====================================
|
|
# TEXT
|
|
# =====================================
|
|
|
|
try:
|
|
|
|
text = page.extract_text()
|
|
|
|
if text:
|
|
parts.append(text)
|
|
|
|
except Exception as e:
|
|
|
|
logger.warning(
|
|
f"PDF text extraction failed: {e}"
|
|
)
|
|
|
|
# =====================================
|
|
# TABLES
|
|
# =====================================
|
|
|
|
try:
|
|
|
|
tables = page.extract_tables()
|
|
|
|
for table_idx, table in enumerate(
|
|
tables,
|
|
start=1
|
|
):
|
|
|
|
parts.append(
|
|
f"=== TABLE {table_idx} ==="
|
|
)
|
|
|
|
for row in table:
|
|
|
|
if not row:
|
|
continue
|
|
|
|
cells = []
|
|
|
|
for cell in row:
|
|
|
|
if cell is None:
|
|
continue
|
|
|
|
txt = str(cell).strip()
|
|
|
|
txt = " ".join(
|
|
txt.split()
|
|
)
|
|
|
|
if txt:
|
|
cells.append(txt)
|
|
|
|
if cells:
|
|
|
|
parts.append(
|
|
" | ".join(cells)
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
logger.warning(
|
|
f"PDF table extraction failed: {e}"
|
|
)
|
|
|
|
text = "\n".join(parts)
|
|
|
|
return self._cleanup_text(text)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(f"PDF parse failed: {e}")
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# DOCX / DOCM
|
|
# =========================================================
|
|
|
|
def _extract_docx(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
doc = Document(io.BytesIO(content))
|
|
|
|
parts = []
|
|
|
|
# =============================================
|
|
# PARAGRAPHS
|
|
# =============================================
|
|
|
|
for p in doc.paragraphs:
|
|
|
|
txt = p.text.strip()
|
|
|
|
if txt:
|
|
parts.append(txt)
|
|
|
|
# =============================================
|
|
# TABLES
|
|
# =============================================
|
|
|
|
for table_idx, table in enumerate(
|
|
doc.tables
|
|
):
|
|
|
|
parts.append(
|
|
f"=== TABLE {table_idx + 1} ==="
|
|
)
|
|
|
|
for row in table.rows:
|
|
|
|
cells = []
|
|
|
|
for cell in row.cells:
|
|
|
|
cell_parts = []
|
|
|
|
# =================================
|
|
# DIRECT CELL TEXT
|
|
# =================================
|
|
|
|
txt = "\n".join(
|
|
p.text.strip()
|
|
for p in cell.paragraphs
|
|
if p.text.strip()
|
|
)
|
|
|
|
txt = " ".join(txt.split())
|
|
|
|
if txt:
|
|
cell_parts.append(txt)
|
|
|
|
# =================================
|
|
# NESTED TABLES
|
|
# =================================
|
|
|
|
for nested_table in cell.tables:
|
|
|
|
for nested_row in nested_table.rows:
|
|
|
|
nested_vals = []
|
|
|
|
for nested_cell in nested_row.cells:
|
|
|
|
nested_txt = "\n".join(
|
|
p.text.strip()
|
|
for p in nested_cell.paragraphs
|
|
if p.text.strip()
|
|
)
|
|
|
|
nested_txt = " ".join(
|
|
nested_txt.split()
|
|
)
|
|
|
|
if nested_txt:
|
|
nested_vals.append(
|
|
nested_txt
|
|
)
|
|
|
|
if nested_vals:
|
|
|
|
cell_parts.append(
|
|
" || ".join(
|
|
nested_vals
|
|
)
|
|
)
|
|
|
|
# =================================
|
|
# FINAL CELL
|
|
# =================================
|
|
|
|
if cell_parts:
|
|
|
|
cells.append(
|
|
" <<<TABLE>>> ".join(
|
|
cell_parts
|
|
)
|
|
)
|
|
|
|
if cells:
|
|
|
|
parts.append(
|
|
" | ".join(cells)
|
|
)
|
|
|
|
text = "\n".join(parts)
|
|
|
|
return self._cleanup_text(text)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(f"DOCX parse failed: {e}")
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# DOC
|
|
# =========================================================
|
|
|
|
def _extract_doc(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
|
|
input_path = os.path.join(
|
|
tmpdir,
|
|
"temp.doc"
|
|
)
|
|
|
|
with open(input_path, "wb") as f:
|
|
f.write(content)
|
|
|
|
subprocess.run(
|
|
[
|
|
"libreoffice",
|
|
"--headless",
|
|
"--convert-to",
|
|
"docx",
|
|
input_path,
|
|
"--outdir",
|
|
tmpdir,
|
|
],
|
|
check=True,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
|
|
converted_path = os.path.join(
|
|
tmpdir,
|
|
"temp.docx"
|
|
)
|
|
|
|
if not os.path.exists(
|
|
converted_path
|
|
):
|
|
|
|
logger.error(
|
|
"DOC conversion failed"
|
|
)
|
|
|
|
return ""
|
|
|
|
with open(
|
|
converted_path,
|
|
"rb"
|
|
) as f:
|
|
|
|
converted_content = f.read()
|
|
|
|
return self._extract_docx(
|
|
converted_content
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(f"DOC parse failed: {e}")
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# RTF
|
|
# =========================================================
|
|
|
|
def _extract_rtf(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
try:
|
|
text = content.decode("utf-8")
|
|
|
|
except Exception:
|
|
text = content.decode(
|
|
"cp1251",
|
|
errors="ignore"
|
|
)
|
|
|
|
text = rtf_to_text(text)
|
|
|
|
return self._cleanup_text(text)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(f"RTF parse failed: {e}")
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# XLSX
|
|
# =========================================================
|
|
|
|
def _extract_excel_openpyxl(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
wb = openpyxl.load_workbook(
|
|
io.BytesIO(content),
|
|
data_only=True
|
|
)
|
|
|
|
text = []
|
|
|
|
for sheet in wb.worksheets:
|
|
|
|
text.append(
|
|
f"=== SHEET: {sheet.title} ==="
|
|
)
|
|
|
|
# =====================================
|
|
# MERGED CELLS
|
|
# =====================================
|
|
|
|
merged_ranges = list(
|
|
sheet.merged_cells.ranges
|
|
)
|
|
|
|
if merged_ranges:
|
|
|
|
text.append(
|
|
"=== MERGED CELLS ==="
|
|
)
|
|
|
|
for r in merged_ranges:
|
|
text.append(str(r))
|
|
|
|
# =====================================
|
|
# ROWS
|
|
# =====================================
|
|
|
|
text.append("[ROWS]")
|
|
|
|
for row in sheet.iter_rows(
|
|
values_only=True
|
|
):
|
|
|
|
vals = [
|
|
self._cell_to_text(c)
|
|
for c in row
|
|
]
|
|
|
|
vals = [x for x in vals if x]
|
|
|
|
if vals:
|
|
|
|
text.append(
|
|
" | ".join(vals)
|
|
)
|
|
|
|
# =====================================
|
|
# COLUMNS
|
|
# =====================================
|
|
|
|
text.append("[COLUMNS]")
|
|
|
|
max_col = sheet.max_column or 0
|
|
max_row = sheet.max_row or 0
|
|
|
|
for c in range(1, max_col + 1):
|
|
|
|
vals = []
|
|
|
|
for r in range(1, max_row + 1):
|
|
|
|
v = self._cell_to_text(
|
|
sheet.cell(r, c).value
|
|
)
|
|
|
|
if v:
|
|
vals.append(v)
|
|
|
|
if vals:
|
|
|
|
col_letter = (
|
|
openpyxl.utils
|
|
.get_column_letter(c)
|
|
)
|
|
|
|
text.append(
|
|
f"COL {col_letter}: "
|
|
+ " || ".join(vals)
|
|
)
|
|
|
|
return self._cleanup_text(
|
|
"\n".join(text)
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(
|
|
f"XLSX parse failed: {e}"
|
|
)
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# XLS
|
|
# =========================================================
|
|
|
|
def _extract_excel_xls(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
import xlrd
|
|
|
|
workbook = xlrd.open_workbook(
|
|
file_contents=content
|
|
)
|
|
|
|
text = []
|
|
|
|
for sheet in workbook.sheets():
|
|
|
|
text.append(
|
|
f"=== SHEET: {sheet.name} ==="
|
|
)
|
|
|
|
for row_idx in range(
|
|
sheet.nrows
|
|
):
|
|
|
|
row = sheet.row_values(
|
|
row_idx
|
|
)
|
|
|
|
vals = [
|
|
self._cell_to_text(c)
|
|
for c in row
|
|
]
|
|
|
|
vals = [x for x in vals if x]
|
|
|
|
if vals:
|
|
|
|
text.append(
|
|
" | ".join(vals)
|
|
)
|
|
|
|
return self._cleanup_text(
|
|
"\n".join(text)
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(f"XLS parse failed: {e}")
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# CSV
|
|
# =========================================================
|
|
|
|
def _extract_csv(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
try:
|
|
text = content.decode("utf-8")
|
|
|
|
except Exception:
|
|
text = content.decode(
|
|
"cp1251",
|
|
errors="ignore"
|
|
)
|
|
|
|
reader = csv.reader(
|
|
io.StringIO(text)
|
|
)
|
|
|
|
rows = []
|
|
|
|
for row in reader:
|
|
|
|
vals = [
|
|
str(x).strip()
|
|
for x in row
|
|
if str(x).strip()
|
|
]
|
|
|
|
if vals:
|
|
|
|
rows.append(
|
|
" | ".join(vals)
|
|
)
|
|
|
|
return self._cleanup_text(
|
|
"\n".join(rows)
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(f"CSV parse failed: {e}")
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# OCR
|
|
# =========================================================
|
|
|
|
def _ocr_image(
|
|
self,
|
|
image: Image.Image
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
text = pytesseract.image_to_string(
|
|
image,
|
|
lang="rus+eng",
|
|
config="--psm 6"
|
|
)
|
|
|
|
return self._cleanup_text(text)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(f"OCR failed: {e}")
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# IMAGE
|
|
# =========================================================
|
|
|
|
def _extract_image(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
image = Image.open(
|
|
io.BytesIO(content)
|
|
)
|
|
|
|
return self._ocr_image(image)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(
|
|
f"Image parse failed: {e}"
|
|
)
|
|
|
|
return ""
|
|
|
|
# =========================================================
|
|
# TIFF
|
|
# =========================================================
|
|
|
|
def _extract_tiff(
|
|
self,
|
|
content: bytes
|
|
) -> str:
|
|
|
|
try:
|
|
|
|
image = Image.open(
|
|
io.BytesIO(content)
|
|
)
|
|
|
|
texts = []
|
|
|
|
for frame in range(
|
|
getattr(image, "n_frames", 1)
|
|
):
|
|
|
|
image.seek(frame)
|
|
|
|
texts.append(
|
|
self._ocr_image(image)
|
|
)
|
|
|
|
return self._cleanup_text(
|
|
"\n".join(texts)
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
logger.error(
|
|
f"TIFF parse failed: {e}"
|
|
)
|
|
|
|
return ""
|