接口自动化之excel读写封装

本次封装基于openpyxl进行的二次封装

安装openpyxl

pip install openpyxl

封装ExcelReader

采用yield的方式返回数据，减少内存的占用

class ExcelReader:

    def __init__(self, filename):
        self._excel: Workbook = load_workbook(filename, read_only=True)

    def read(self, sheet=None):
        if not sheet:
            sheet = self._excel.active
        else:
            sheet = self._excel[sheet]
        rows = sheet.rows
        title = [str(cell.value) for cell in next(rows)]
        for row in rows:
            yield dict(zip(title, (cell.value for cell in row)))

封装ExcelWriter

class ExcelWriter:
    def __init__(self, filename):
        self._file = filename
        self._excel = Workbook(write_only=True)
        self._title_writer_flag = False
        self._value_order = []
        self._title_width = []
        self._title_font = Font(name='等线', size=12, bold=True)
        self._title_fill = PatternFill('solid', fgColor='95b3d7')
        self._font = Font(name='等线', size=12)

    def _get_title_cell(self, sheet, value):
        cell = WriteOnlyCell(sheet, value=value)
        cell.font = self._title_font
        cell.fill = self._title_fill
        return cell

    def _write_title(self, sheet, title):
        self._value_order = title  # 定义列顺序
        sheet.append((self._get_title_cell(sheet, t) for t in self._value_order))
        self._title_writer_flag = True

    def _get_cell(self, sheet, value):
        cell = WriteOnlyCell(sheet, value=value)
        cell.font = self._font
        return cell

    def _write_row(self, sheet, row):
        sheet.append((self._get_cell(sheet, val) for val in row))

    @staticmethod
    def _set_column_width(sheet, title):
        for i, t in enumerate(title, start=1):
            column_title = get_column_letter(i)
            width = 10  # 定义每一列的宽度，此处可以按照数据来设置列宽
            sheet.column_dimensions[column_title].width = width

    def write(self, sheet_name: str, data, title: Iterable[str] | None = None):
        sheet = self._excel.create_sheet(sheet_name)
        if title:
            self._write_title(sheet, title)
            self._set_column_width(sheet, title)
        for row in data:
            if isinstance(row, dict):
                if not self._title_writer_flag:
                    self._write_title(sheet, row.keys())
                row_value = (row.get(key) for key in self._value_order)
            else:
                row_value = row
            sheet.append((self._get_cell(sheet, val) for val in row_value))
        return sheet

    def save(self):
        self._excel.save(self._file)