1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# Copyright (c) 2024 Huawei Device Co., Ltd.
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import xlwt
17from xlwt import Worksheet
18from xlwt import Pattern as PTRN, Alignment as ALMT, Font as FT, XFStyle as XFSTL
19import typing
20from typing import Optional
21from collections.abc import Iterable
22
23
24class SimpleExcelWriter:
25    def __init__(self, default_sheet_name: str = "sheet1"):
26        self.__book = xlwt.Workbook(encoding='utf-8', style_compression=0)
27        self.__sheet_dict = {
28            default_sheet_name: self.__book.add_sheet(
29                sheetname=default_sheet_name, cell_overwrite_ok=True)
30        }
31        self.__sheet_pos = {
32            default_sheet_name: (0, 0)  # 记录各个sheet页已经写到什么位置了，当前值为还没有写的
33        }
34        self.__default_sheet_name = default_sheet_name
35        # 表头样式
36        self.__head_style = XFSTL()
37        # 内容样式
38        self.__content_style = XFSTL()
39        # 字体
40        ft = FT()
41        ft.bold = True
42
43        # 设置背景颜色
44        ptrn = PTRN()
45        ptrn.pattern = PTRN.SOLID_PATTERN
46        ptrn.pattern_fore_colour = 22  # 背景颜色
47
48        # 居中对齐
49        algmt = ALMT()
50        algmt.horz = ALMT.HORZ_CENTER  # 水平方向
51        algmt.vert = ALMT.VERT_CENTER  # 垂直方向
52
53        self.__head_style.font = ft
54        self.__head_style.alignment = algmt
55        self.__head_style.pattern = ptrn
56        self.__content_style.alignment = algmt
57
58    def append_line(self, content: list, sheet_name: str = None):
59        sheet_name = self.__default_sheet_name if sheet_name is None else sheet_name
60        if sheet_name not in self.__sheet_dict.keys():
61            print("error: sheet name '{}' not exist".format(sheet_name))
62            return
63        sheet: Worksheet = self.__sheet_dict.get(sheet_name)
64        x, y = self.__sheet_pos.get(sheet_name)
65        for ele in content:
66            sheet.write(x, y, ele, style=self.__content_style)
67            y = self.__increment_y(sheet_name)
68        self.__increment_x(sheet_name)
69
70    def write_merge(self, x0: int, y0: int, x1: int, y1: int, content: typing.Any,
71                    sheet_name: str = None):
72        sheet_name = self.__default_sheet_name if sheet_name is None else sheet_name
73        if sheet_name not in self.__sheet_dict.keys():
74            print("error: sheet name '{}' not exist".format(sheet_name))
75            return
76        sheet: Worksheet = self.__sheet_dict.get(sheet_name)
77        sheet.write_merge(x0, x1, y0, y1, content, style=self.__content_style)
78
79    def set_sheet_header(self, headers: Iterable, sheet_name: str = None):
80        """
81        给sheet页设置表头
82        """
83        sheet_name = self.__default_sheet_name if sheet_name is None else sheet_name
84        if sheet_name not in self.__sheet_dict.keys():
85            print("error: sheet name '{}' not exist".format(sheet_name))
86            return
87        x, y = self.__sheet_pos.get(sheet_name)
88        if x != 0 or y != 0:
89            print(
90                "error: pos of sheet '{}' is not (0,0). set_sheet_header must before write".format(sheet_name))
91            return
92        sheet: Worksheet = self.__sheet_dict.get(sheet_name)
93        for h in headers:
94            sheet.write(x, y, h, self.__head_style)
95            y = self.__increment_y(sheet_name)
96        self.__increment_x(sheet_name)
97
98    def add_sheet(self, sheet_name: str, cell_overwrite_ok=True) -> Optional[xlwt.Worksheet]:
99        if sheet_name in self.__sheet_dict.keys():
100            print("error: sheet name '{}' has exist".format(sheet_name))
101            return
102        self.__sheet_dict[sheet_name] = self.__book.add_sheet(
103            sheetname=sheet_name, cell_overwrite_ok=cell_overwrite_ok)
104        self.__sheet_pos[sheet_name] = (0, 0)
105        return self.__sheet_dict.get(sheet_name)
106
107    def save(self, file_name: str):
108        self.__book.save(file_name)
109
110    def __increment_y(self, sheet_name: str, value: int = 1) -> int:
111        if sheet_name in self.__sheet_pos.keys():
112            x, y = self.__sheet_pos.get(sheet_name)
113            y = y + value
114            self.__sheet_pos[sheet_name] = (x, y)
115            return y
116
117    def __increment_x(self, sheet_name: str, value: int = 1) -> int:
118        if sheet_name in self.__sheet_pos.keys():
119            x, y = self.__sheet_pos.get(sheet_name)
120            x = x + value
121            self.__sheet_pos[sheet_name] = (x, 0)
122            return x
123
124
125if __name__ == '__main__':
126    writer = SimpleExcelWriter(default_sheet_name="first")
127    writer.add_sheet("second")
128    writer.add_sheet("third")
129    writer.set_sheet_header(["h", "m", "n"])
130    writer.append_line([1, 2, 3])
131    writer.append_line([2, 3, 4], "second")
132    writer.append_line([3, 4, 5], "third")
133    writer.append_line([3, 2, 1])
134    writer.save("demo.xls")
135