easy-ml/monitor/report_generate.py

# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/11/8
@desc: 
"""
import os
from typing import Dict

from docx import Document
from docx.enum.table import WD_ALIGN_VERTICAL
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.shared import Inches, Cm

from commom import GeneralException, f_get_datetime
from config import BaseConfig
from entitys import MetricFucEntity
from enums import ResultCodesEnum, PlaceholderPrefixEnum


class Report():

    @staticmethod
    def _set_cell_width(cell):
        text = cell.text
        if len(text) >= 10:
            cell.width = Cm(2)
        elif len(text) >= 15:
            cell.width = Cm(2.5)
        elif len(text) >= 25:
            cell.width = Cm(3)
        else:
            cell.width = Cm(1.5)

    @staticmethod
    def _set_cell_format(cell):
        cell.paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
        cell.vertical_alignment = WD_ALIGN_VERTICAL.CENTER

    @staticmethod
    def _merge_cell_column(pre_cell, curr_cell):
        if curr_cell.text == pre_cell.text:
            column_name = curr_cell.text
            pre_cell.merge(curr_cell)
            pre_cell.text = column_name
            for run in pre_cell.paragraphs[0].runs:
                run.bold = True
            Report._set_cell_format(pre_cell)
            Report._set_cell_width(pre_cell)

    @staticmethod
    def _set_table_singleBoard(table):
        # 将table 的所有单元格四个边设置为 0.5 镑, 黑色, 实线

        def _set_table_boarder(table, **kwargs):
            """
            Set table`s border
            Usage:
            set_table_border(
                cell,
                top={"sz": 12, "val": "single", "color": "#FF0000"},
                bottom={"sz": 12, "color": "#00FF00", "val": "single"},
                left={"sz": 24, "val": "dashed"},
                right={"sz": 12, "val": "dashed"},
            )
            """
            borders = OxmlElement('w:tblBorders')
            for tag in ('bottom', 'top', 'left', 'right', 'insideV', 'insideH'):
                edge_data = kwargs.get(tag)
                if edge_data:
                    any_border = OxmlElement(f'w:{tag}')
                    for key in ["sz", "val", "color", "space", "shadow"]:
                        if key in edge_data:
                            any_border.set(qn(f'w:{key}'), str(edge_data[key]))
                    borders.append(any_border)
                    table._tbl.tblPr.append(borders)

        return _set_table_boarder(
            table,
            top={"sz": 4, "val": "single", "color": "#000000"},
            bottom={"sz": 4, "val": "single", "color": "#000000"},
            left={"sz": 4, "val": "single", "color": "#000000"},
            right={"sz": 4, "val": "single", "color": "#000000"},
            insideV={"sz": 4, "val": "single", "color": "#000000"},
            insideH={"sz": 4, "val": "single", "color": "#000000"}
        )

    @staticmethod
    def _get_placeholder(placeholder_prefix_enum: PlaceholderPrefixEnum, metric_code: str):
        return "{{" + f"{placeholder_prefix_enum.value}{metric_code}" + "}}"

    @staticmethod
    def _fill_value_placeholder(doc: Document, metric_value_dict: Dict[str, MetricFucEntity]):
        # 替换指标
        for paragraph in doc.paragraphs:
            text = paragraph.text
            for metric_code, metric_fuc_entity in metric_value_dict.items():
                placeholder = Report._get_placeholder(PlaceholderPrefixEnum.VALUE, metric_code)
                metric_value = metric_fuc_entity.value
                if metric_value is None:
                    continue
                text = text.replace(placeholder, metric_value)
            # 段落中多个runs时执行，最后一个run改成替换好的文本，其他run置空
            if len(paragraph.runs[:-1]) > 0:
                for run in paragraph.runs[:-1]:
                    run.text = ''
                paragraph.runs[-1].text = text

    @staticmethod
    def _fill_table_placeholder(doc: Document, metric_value_dict: Dict[str, MetricFucEntity]):
        # 替换表格
        for paragraph in doc.paragraphs:
            for metric_code, metric_fuc_entity in metric_value_dict.items():
                placeholder = Report._get_placeholder(PlaceholderPrefixEnum.TABLE, metric_code)
                metric_table = metric_fuc_entity.table
                if metric_table is None:
                    continue
                if not placeholder in paragraph.text:
                    continue
                # 清除占位符
                for run in paragraph.runs:
                    run.text = run.text.replace(placeholder, "")
                table = doc.add_table(rows=metric_table.shape[0] + 1, cols=metric_table.shape[1])
                table.alignment = WD_ALIGN_PARAGRAPH.CENTER
                paragraph._element.addnext(table._element)
                # 列名
                for column_idx, column_name in enumerate(metric_table.columns):
                    cell = table.cell(0, column_idx)
                    cell.text = str(column_name)
                    for run in cell.paragraphs[0].runs:
                        run.bold = True
                    Report._set_cell_format(cell)
                    Report._set_cell_width(cell)
                    # 合并相同的列名
                    if column_idx != 0 and BaseConfig.merge_table_column:
                        pre_cell = table.cell(0, column_idx - 1)
                        Report._merge_cell_column(pre_cell, cell)
                # 值
                for row_idx, row in metric_table.iterrows():
                    for column_idx, value in enumerate(row):
                        cell = table.cell(row_idx + 1, column_idx)
                        cell.text = str(value)
                        Report._set_cell_format(cell)
                        Report._set_cell_width(cell)
                        # 合并第一行数据也为列的情况
                        if row_idx == 0:
                            Report._merge_cell_column(table.cell(0, column_idx), cell)

                Report._set_table_singleBoard(table)
                # 禁止自动调整表格
                if len(metric_table.columns) <= 12:
                    table.autofit = False

    @staticmethod
    def _fill_image_placeholder(doc: Document, metric_value_dict: Dict[str, MetricFucEntity]):
        # 替换图片
        for paragraph in doc.paragraphs:
            for metric_code, metric_fuc_entity in metric_value_dict.items():
                placeholder = Report._get_placeholder(PlaceholderPrefixEnum.IMAGE, metric_code)
                image_path = metric_fuc_entity.image_path
                if image_path is None:
                    continue
                if not placeholder in paragraph.text:
                    continue
                if not os.path.exists(image_path):
                    raise GeneralException(ResultCodesEnum.NOT_FOUND, message=f"文件【{image_path}】不存在")
                # 清除占位符
                for run in paragraph.runs:
                    if placeholder not in run.text:
                        continue
                    run.text = run.text.replace(placeholder, "")
                    run.add_picture(image_path, width=Inches(6))

    @staticmethod
    def generate_report(metric_value_dict: Dict[str, MetricFucEntity], template_path: str):
        if os.path.exists(template_path):
            doc = Document(template_path)
        else:
            raise GeneralException(ResultCodesEnum.NOT_FOUND, message=f"监控模板文件【{template_path}】不存在")

        Report._fill_value_placeholder(doc, metric_value_dict)
        Report._fill_table_placeholder(doc, metric_value_dict)
        Report._fill_image_placeholder(doc, metric_value_dict)
        new_path = template_path.replace(".docx", f"{f_get_datetime()}.docx")
        doc.save(f"./{new_path}")


if __name__ == "__main__":
    pass