"""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