285 lines
12 KiB
Python
285 lines
12 KiB
Python
"""ESPD DOCX generator (ported from GostGenerator2)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
from docx.enum.table import WD_CELL_VERTICAL_ALIGNMENT
|
|
from docx.enum.text import WD_ALIGN_PARAGRAPH
|
|
from docx.oxml import OxmlElement
|
|
from docx.oxml.ns import qn
|
|
from docx.shared import Pt
|
|
from docx.table import Table
|
|
from docxtpl import DocxTemplate
|
|
|
|
from app.services.espd.models import (
|
|
CONTENT_CONTEXT_KEYS,
|
|
DOC_TYPE_CONFIG,
|
|
CommonFields,
|
|
DocumentType,
|
|
GenerateRequest,
|
|
ImpZipInfo,
|
|
SelectedDocument,
|
|
SpecItem,
|
|
)
|
|
from app.services.espd.html_util import html_to_plain
|
|
|
|
|
|
TEMPLATES_DIR = Path(__file__).resolve().parents[2] / "espd_templates"
|
|
|
|
|
|
class DocumentGenerator:
|
|
def __init__(self, templates_dir: str | Path | None = None):
|
|
self.templates_dir = Path(templates_dir or TEMPLATES_DIR).resolve()
|
|
|
|
def get_template_path(self, doc_type: DocumentType) -> Path:
|
|
config = DOC_TYPE_CONFIG.get(doc_type)
|
|
if not config:
|
|
raise ValueError(f"Неизвестный тип документа: {doc_type}")
|
|
return self.templates_dir / config["template"]
|
|
|
|
def build_context(self, request: GenerateRequest) -> dict:
|
|
config = DOC_TYPE_CONFIG.get(request.doc_type)
|
|
context = {
|
|
"DECIMAL": request.common.decimal,
|
|
"PROGNAME": request.common.progname,
|
|
"AUTHOR": request.common.author,
|
|
"CHECKER": request.common.checker,
|
|
"NCONTROL": request.common.ncontrol,
|
|
"APPROVER": request.common.approver,
|
|
"YEAR": request.common.year,
|
|
"decimal": request.common.decimal,
|
|
"progname": request.common.progname,
|
|
"author": request.common.author,
|
|
"checker": request.common.checker,
|
|
"ncontrol": request.common.ncontrol,
|
|
"approver": request.common.approver,
|
|
"year": request.common.year,
|
|
"doc_type_name": config["name"] if config else "",
|
|
"doc_type_code": config["code"] if config else "",
|
|
"doc_suffix": config["suffix"] if config else "",
|
|
"DECIMAL_FULL": request.common.decimal + (config["suffix"] if config else ""),
|
|
"SPEC_TABLE_PLACEHOLDER": "",
|
|
"RO_CONTENT": (request.extra_context or {}).get("RO_CONTENT", ""),
|
|
"DOC_CONTENT": (request.extra_context or {}).get("DOC_CONTENT", ""),
|
|
"TEXT_CONTENT": (request.extra_context or {}).get("TEXT_CONTENT", ""),
|
|
"RSP_CONTENT": (request.extra_context or {}).get("RSP_CONTENT", ""),
|
|
"RP_CONTENT": (request.extra_context or {}).get("RP_CONTENT", ""),
|
|
"PMI_CONTENT": (request.extra_context or {}).get("PMI_CONTENT", ""),
|
|
"IMP_CONTENT": (request.extra_context or {}).get("IMP_CONTENT", ""),
|
|
}
|
|
if request.imp_zip_info:
|
|
context["FOLDER_TABLE_PLACEHOLDER"] = "%%FOLDER_TABLE%%"
|
|
context["ZIP_CHECKSUM"] = request.imp_zip_info.zip_checksum
|
|
else:
|
|
context["FOLDER_TABLE_PLACEHOLDER"] = ""
|
|
context["ZIP_CHECKSUM"] = ""
|
|
if request.extra_context:
|
|
context.update(request.extra_context)
|
|
return context
|
|
|
|
def set_row_height(self, row, height_cm: float) -> None:
|
|
tr = row._tr
|
|
tr_pr = tr.get_or_add_trPr()
|
|
tr_height = OxmlElement("w:trHeight")
|
|
tr_height.set(qn("w:val"), str(int(height_cm * 567)))
|
|
tr_height.set(qn("w:hRule"), "exact")
|
|
tr_pr.append(tr_height)
|
|
|
|
def format_cell(self, cell, font_size: int = 14) -> None:
|
|
cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER
|
|
for paragraph in cell.paragraphs:
|
|
paragraph.alignment = WD_ALIGN_PARAGRAPH.LEFT
|
|
for run in paragraph.runs:
|
|
run.font.size = Pt(font_size)
|
|
|
|
def add_empty_row(self, table, row_height_cm: float = 0.7):
|
|
row = table.add_row()
|
|
self.set_row_height(row, row_height_cm)
|
|
for cell in row.cells:
|
|
cell.text = ""
|
|
self.format_cell(cell, font_size=14)
|
|
return row
|
|
|
|
def add_spec_rows_to_table(self, table: Table, spec_items: list, rows_per_page: int = 35) -> None:
|
|
row_count = 0
|
|
for item in spec_items:
|
|
row = table.add_row()
|
|
row_count += 1
|
|
self.set_row_height(row, 0.7)
|
|
row.cells[0].text = item.decimal
|
|
row.cells[1].text = item.name
|
|
row.cells[2].text = item.note
|
|
for cell in row.cells:
|
|
self.format_cell(cell, font_size=14)
|
|
total_rows_needed = rows_per_page - 1
|
|
while row_count < total_rows_needed:
|
|
self.add_empty_row(table, 0.7)
|
|
row_count += 1
|
|
|
|
def _format_table_cell_font(self, cell, font_size: int = 14) -> None:
|
|
for paragraph in cell.paragraphs:
|
|
for run in paragraph.runs:
|
|
run.font.size = Pt(font_size)
|
|
|
|
def _get_cell_text(self, cell) -> str:
|
|
text = "".join(p.text or "" for p in cell.paragraphs)
|
|
if text.strip():
|
|
return text
|
|
return "".join((node.text or "") for node in cell._tc.iter() if node.tag == qn("w:t"))
|
|
|
|
def _replace_folder_placeholder_with_table(
|
|
self, doc: DocxTemplate, folder_table_rows: list[list[str]]
|
|
) -> None:
|
|
for table in doc.docx.tables:
|
|
if len(table.rows) < 1 or len(table.columns) < 2:
|
|
continue
|
|
first_row = table.rows[0]
|
|
c0 = self._get_cell_text(first_row.cells[0]).strip()
|
|
c1 = self._get_cell_text(first_row.cells[1]).strip()
|
|
if "Имя файла" not in c0 or "Размер" not in c1:
|
|
continue
|
|
rows_to_remove = []
|
|
for row_idx, row in enumerate(table.rows):
|
|
if row_idx == 0:
|
|
continue
|
|
for cell in row.cells:
|
|
if "FOLDER_TABLE" in self._get_cell_text(cell) or "%%" in self._get_cell_text(cell):
|
|
rows_to_remove.append(row._tr)
|
|
break
|
|
for tr in reversed(rows_to_remove):
|
|
table._tbl.remove(tr)
|
|
for name, size_str in folder_table_rows:
|
|
new_row = table.add_row()
|
|
new_row.cells[0].text = name
|
|
new_row.cells[1].text = size_str
|
|
self._format_table_cell_font(new_row.cells[0], 14)
|
|
self._format_table_cell_font(new_row.cells[1], 14)
|
|
for row in table.rows:
|
|
for cell in row.cells:
|
|
self._format_table_cell_font(cell, 14)
|
|
return
|
|
|
|
def generate(self, request: GenerateRequest) -> io.BytesIO:
|
|
template_path = self.get_template_path(request.doc_type)
|
|
if not template_path.exists():
|
|
raise FileNotFoundError(f"Шаблон не найден: {template_path}")
|
|
doc = DocxTemplate(template_path)
|
|
context = self.build_context(request)
|
|
doc.render(context)
|
|
if request.doc_type == DocumentType.SPEC and request.spec_items:
|
|
if doc.docx.tables:
|
|
self.add_spec_rows_to_table(doc.docx.tables[0], request.spec_items)
|
|
for para in doc.docx.paragraphs:
|
|
if "SPEC_TABLE_PLACEHOLDER" in para.text:
|
|
p = para._element
|
|
p.getparent().remove(p)
|
|
break
|
|
if request.doc_type == DocumentType.IMP and request.imp_zip_info:
|
|
self._replace_folder_placeholder_with_table(doc, request.imp_zip_info.folder_table_rows)
|
|
output = io.BytesIO()
|
|
doc.save(output)
|
|
output.seek(0)
|
|
return output
|
|
|
|
def generate_filename(self, request: GenerateRequest) -> str:
|
|
config = DOC_TYPE_CONFIG.get(request.doc_type, {})
|
|
suffix = config.get("suffix", "")
|
|
safe_decimal = request.common.decimal.replace("/", "_").replace("\\", "_")
|
|
return f"{safe_decimal}{suffix}.docx"
|
|
|
|
def generate_lu(self, common: CommonFields, doc_type: DocumentType) -> io.BytesIO:
|
|
template_path = (self.templates_dir / "lu.docx").resolve()
|
|
if not template_path.exists():
|
|
raise FileNotFoundError(f"Шаблон ЛУ не найден: {template_path}")
|
|
config = DOC_TYPE_CONFIG.get(doc_type, {})
|
|
suffix = config.get("suffix", "")
|
|
context = {
|
|
"DECIMAL": common.decimal,
|
|
"PROGNAME": common.progname,
|
|
"AUTHOR": common.author,
|
|
"CHECKER": common.checker,
|
|
"NCONTROL": common.ncontrol,
|
|
"APPROVER": common.approver,
|
|
"YEAR": common.year,
|
|
"DOCTYPE": config.get("name", ""),
|
|
"DOCCODE": config.get("code", ""),
|
|
"DECIMAL_FULL": common.decimal + suffix,
|
|
"DECIMAL_LU": common.decimal + suffix + "-ЛУ",
|
|
}
|
|
doc = DocxTemplate(template_path)
|
|
doc.render(context)
|
|
output = io.BytesIO()
|
|
doc.save(output)
|
|
output.seek(0)
|
|
return output
|
|
|
|
def generate_lu_filename(self, common: CommonFields, doc_type: DocumentType) -> str:
|
|
config = DOC_TYPE_CONFIG.get(doc_type, {})
|
|
suffix = config.get("suffix", "")
|
|
safe_decimal = common.decimal.replace("/", "_").replace("\\", "_")
|
|
return f"{safe_decimal}{suffix}-ЛУ.docx"
|
|
|
|
def generate_all_zip(
|
|
self,
|
|
common: CommonFields,
|
|
selected_docs: list[SelectedDocument],
|
|
spec_items: list[SpecItem],
|
|
imp_zip_info: ImpZipInfo | None = None,
|
|
ro_ai_content: str | None = None,
|
|
doc_contents: dict[str, str] | None = None,
|
|
) -> io.BytesIO:
|
|
zip_buffer = io.BytesIO()
|
|
contents = dict(doc_contents or {})
|
|
if ro_ai_content and "ro" not in contents:
|
|
contents["ro"] = ro_ai_content
|
|
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
|
|
for selected in selected_docs:
|
|
doc_type = selected.doc_type
|
|
if doc_type == DocumentType.SPEC:
|
|
lu_content = self.generate_lu(common, doc_type)
|
|
zip_file.writestr(self.generate_lu_filename(common, doc_type), lu_content.read())
|
|
else:
|
|
template_path = self.get_template_path(doc_type)
|
|
if template_path.exists():
|
|
plain = contents.get(doc_type.value, "")
|
|
if plain and "<" in plain:
|
|
plain = html_to_plain(plain)
|
|
key = CONTENT_CONTEXT_KEYS.get(doc_type)
|
|
extra = {"DOC_CONTENT": plain}
|
|
if key:
|
|
extra[key] = plain
|
|
doc_request = GenerateRequest(
|
|
doc_type=doc_type,
|
|
common=common,
|
|
spec_items=[],
|
|
imp_zip_info=imp_zip_info if doc_type == DocumentType.IMP else None,
|
|
extra_context=extra,
|
|
)
|
|
doc_content = self.generate(doc_request)
|
|
zip_file.writestr(self.generate_filename(doc_request), doc_content.read())
|
|
lu_content = self.generate_lu(common, doc_type)
|
|
zip_file.writestr(self.generate_lu_filename(common, doc_type), lu_content.read())
|
|
if spec_items:
|
|
spec_request = GenerateRequest(
|
|
doc_type=DocumentType.SPEC,
|
|
common=common,
|
|
spec_items=spec_items,
|
|
)
|
|
spec_content = self.generate(spec_request)
|
|
zip_file.writestr(self.generate_filename(spec_request), spec_content.read())
|
|
zip_buffer.seek(0)
|
|
return zip_buffer
|
|
|
|
|
|
_generator: DocumentGenerator | None = None
|
|
|
|
|
|
def get_generator() -> DocumentGenerator:
|
|
global _generator
|
|
if _generator is None:
|
|
_generator = DocumentGenerator()
|
|
return _generator
|