Pandas 导出到 Excel 并设置自适应列宽

Posted 2023-03-03 Aspirantlu

主要逻辑在 CreateExcel 和 AddInfoToExcel 方法中。

# cython: language_level=3

import typing  # noqa
import pandas
import os
from pathlib import Path
import collections
from copy import copy
from openpyxl.utils import get_column_letter
import numpy as np

from Nkd.Test.BusinessOrchestration.TestManagement.InputObjects.TestInput import TestInput
from Nkd.Test.BusinessOrchestration.TestManagement.OutputObjects.TestOutput import TestOutput

from Nkd.Test.Common.Utilities import Utilities
from Nkd.Test.BusinessObjects.SAPMaterialInformation import SAPMaterialInformationCollection
from Nkd.Test.Common.NkdException import NkdException


class TestOrchestration(object):

    objectTypeName = 'Nkd.Test.BusinessOrchestration.TestManagement.TestOrchestration'

    @classmethod
    def Test(cls, input: TestInput) -> TestOutput:
        Utilities.StartMethod(cls.objectTypeName, 'Test', 
                              'TestInput': input)
        try:
            """导出SAPMaterialInformation数据"""
            output = TestOutput()
            if not input.FileName:
                raise NkdException("未传入文件名")
            sap_material_information = input.SAPMaterialInformation
            if sap_material_information:
                # 导出传入的数据
                for i in sap_material_information:
                    if not i.ObjectExists():
                        raise NkdException(f"Id为i.Id的数据不存在")
                sap_material_information.Load()
            else:
                # 导出全部
                sap_material_information = SAPMaterialInformationCollection()
                sap_material_information.LoadByFilter()
            file_path = Path("Share/Temp/")
            if not file_path.exists():
                os.makedirs(file_path)
            excel_file = f"Share/Temp/input.FileName.xlsx"

            cls.CreateExcel(excel_file)
            cls.AddInfoToExcel(excel_file, sap_material_information)

            output.FilePath = excel_file
            Utilities.EndMethod(None, None, 
                                'TestInput': input, 'TestOutput': output)
        except Exception as ex:
            raise ex

        return output

    @staticmethod
    def CreateExcel(excel_file_):
        """创建一个名字为excel_file_name_的excel文件，并设置表头"""
        # 如果文件不存在将创建一个新文件，如果文件存在则将其覆盖。注意覆盖时会删除原文件中的数据。
        form_header = [
            '一', '二', '三', '四', '五'
        ]

        # 指定dtype 用来解决pandas版本问题
        df = pandas.DataFrame(columns=form_header, dtype=object)
        df.to_excel(excel_file_, index=False)

    @classmethod
    def AddInfoToExcel(cls, excel_file_, sap_material_information_: SAPMaterialInformationCollection):
        """把数据插入到excel里面"""
        df = pandas.read_excel(excel_file_, engine="openpyxl", index_col=False)
        columns = [
            'One', 'Two', 'Three', 'Four', 'Five'
        ]

        order_dict_origin = collections.OrderedDict().fromkeys(columns, "")
        row_index = len(df) + 1  # 当前excel内容有几行
        for i in sap_material_information_:
            order_dict = copy(order_dict_origin)
            for k, v in i.__dict__.items():
                if (k not in order_dict) or (v is None):
                    continue
                if k == "One":
                    order_dict[k] = "已签收" if v else "待签收"
                    continue
                order_dict[k] = v
            li = list(order_dict.values())
            df.loc[row_index] = li
            row_index += 1

        df.to_excel(excel_file_, index=False)
        # --以下为了设置自适应列宽
        # 计算表头的字符宽度
        column_widths = (df.columns.to_series().apply(
            lambda x: len(x.encode('gbk'))).values)
        # 计算每列的最大字符宽度
        max_widths = (df.astype(str).applymap(
            lambda x: len(x.encode('gbk'))).agg(max).values)
        # 计算整体最大宽度
        widths = np.max([column_widths, max_widths], axis=0)

        writer = pandas.ExcelWriter(excel_file_, engine='openpyxl')
        df.to_excel(writer, sheet_name='Sheet1', index=False)
        worksheet = writer.sheets['Sheet1']

        for i, width in enumerate(widths, 1):
            worksheet.column_dimensions[get_column_letter(i)].width = width + 3
        writer.save()