142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import mimetypes
|
|
from dataclasses import dataclass
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExtractedBlock:
|
|
text: str
|
|
page_start: int | None = None
|
|
page_end: int | None = None
|
|
section_title: str = ""
|
|
section_path: str = ""
|
|
paragraph_index: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExtractedDocument:
|
|
text: str
|
|
mime_type: str
|
|
blocks: list[ExtractedBlock]
|
|
|
|
|
|
TEXT_EXTENSIONS = {".txt", ".md", ".markdown", ".json", ".csv"}
|
|
DOCX_EXTENSIONS = {".docx"}
|
|
PDF_EXTENSIONS = {".pdf"}
|
|
SUPPORTED_EXTENSIONS = TEXT_EXTENSIONS | DOCX_EXTENSIONS | PDF_EXTENSIONS
|
|
|
|
|
|
def supported_extensions_label() -> str:
|
|
return ", ".join(sorted(SUPPORTED_EXTENSIONS))
|
|
|
|
|
|
def _guess_mime_type(filename: str) -> str:
|
|
guessed, _ = mimetypes.guess_type(filename)
|
|
if guessed:
|
|
return guessed
|
|
ext = Path(filename).suffix.lower()
|
|
if ext == ".docx":
|
|
return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
if ext == ".pdf":
|
|
return "application/pdf"
|
|
return "text/plain"
|
|
|
|
|
|
def extract_document(raw_bytes: bytes, filename: str) -> ExtractedDocument:
|
|
ext = Path(filename or "").suffix.lower()
|
|
if ext in TEXT_EXTENSIONS:
|
|
return _extract_text(raw_bytes, filename)
|
|
if ext in DOCX_EXTENSIONS:
|
|
return _extract_docx(raw_bytes, filename)
|
|
if ext in PDF_EXTENSIONS:
|
|
return _extract_pdf(raw_bytes, filename)
|
|
raise ValueError(f"Неподдерживаемый формат документа. Поддерживаются: {supported_extensions_label()}")
|
|
|
|
|
|
def _extract_text(raw_bytes: bytes, filename: str) -> ExtractedDocument:
|
|
text = raw_bytes.decode("utf-8-sig", errors="replace").strip()
|
|
blocks = [
|
|
ExtractedBlock(text=line.strip(), paragraph_index=idx)
|
|
for idx, line in enumerate(text.splitlines())
|
|
if line.strip()
|
|
]
|
|
if not blocks and text:
|
|
blocks = [ExtractedBlock(text=text, paragraph_index=0)]
|
|
return ExtractedDocument(text=text, mime_type=_guess_mime_type(filename), blocks=blocks)
|
|
|
|
|
|
def _extract_docx(raw_bytes: bytes, filename: str) -> ExtractedDocument:
|
|
try:
|
|
from docx import Document as DocxDocument
|
|
except ImportError as exc:
|
|
raise ValueError("Для загрузки DOCX установите зависимость python-docx") from exc
|
|
|
|
document = DocxDocument(BytesIO(raw_bytes))
|
|
blocks: list[ExtractedBlock] = []
|
|
heading_stack: list[str] = []
|
|
|
|
for idx, paragraph in enumerate(document.paragraphs):
|
|
text = paragraph.text.strip()
|
|
if not text:
|
|
continue
|
|
|
|
style_name = (paragraph.style.name if paragraph.style else "").strip()
|
|
if style_name.lower().startswith("heading"):
|
|
level = _heading_level(style_name)
|
|
heading_stack = heading_stack[: max(level - 1, 0)]
|
|
heading_stack.append(text)
|
|
section_title = text
|
|
else:
|
|
section_title = heading_stack[-1] if heading_stack else ""
|
|
|
|
blocks.append(
|
|
ExtractedBlock(
|
|
text=text,
|
|
section_title=section_title,
|
|
section_path=" > ".join(heading_stack),
|
|
paragraph_index=idx,
|
|
)
|
|
)
|
|
|
|
text = "\n\n".join(block.text for block in blocks).strip()
|
|
return ExtractedDocument(text=text, mime_type=_guess_mime_type(filename), blocks=blocks)
|
|
|
|
|
|
def _heading_level(style_name: str) -> int:
|
|
parts = style_name.rsplit(" ", 1)
|
|
if len(parts) == 2 and parts[1].isdigit():
|
|
return max(1, int(parts[1]))
|
|
return 1
|
|
|
|
|
|
def _extract_pdf(raw_bytes: bytes, filename: str) -> ExtractedDocument:
|
|
try:
|
|
from pypdf import PdfReader
|
|
except ImportError as exc:
|
|
raise ValueError("Для загрузки PDF установите зависимость pypdf") from exc
|
|
|
|
reader = PdfReader(BytesIO(raw_bytes))
|
|
blocks: list[ExtractedBlock] = []
|
|
for page_index, page in enumerate(reader.pages, start=1):
|
|
page_text = (page.extract_text() or "").strip()
|
|
if not page_text:
|
|
continue
|
|
paragraphs = [part.strip() for part in page_text.split("\n\n") if part.strip()]
|
|
if not paragraphs:
|
|
paragraphs = [page_text]
|
|
for paragraph_index, paragraph in enumerate(paragraphs):
|
|
blocks.append(
|
|
ExtractedBlock(
|
|
text=paragraph,
|
|
page_start=page_index,
|
|
page_end=page_index,
|
|
paragraph_index=paragraph_index,
|
|
)
|
|
)
|
|
|
|
text = "\n\n".join(block.text for block in blocks).strip()
|
|
return ExtractedDocument(text=text, mime_type=_guess_mime_type(filename), blocks=blocks)
|