openpyxl utils.

xlutils.py

import re
from enum import Enum
from typing import Union, Optional

from openpyxl.cell.cell import Cell, MergedCell
from openpyxl.worksheet.cell_range import CellRange
from openpyxl.worksheet.worksheet import Worksheet


class Direction(Enum):
    """
    方向向量。
    """
    LEFT = (0, -1)
    UP = (-1, 0)
    RIGHT = (0, 1)
    DOWN = (1, 0)


# Excel 概念：
# Cell：表示普通单元格。
# MergedCell：表示合并单元格（不包括 start_cell），MergedCell 的内容通常总是为 None。
#   如合并单元格（A1:D3），则 A1（即：start_cell） 存储着该合并单元格的值，并且它是 Cell 类型，其余单元格的值为 None，并且他们是 MergedCell 类型。
# start_cell：表示合并单元格的左上角单元格（它也是 Cell 类型），它是 CellRange 的成员。
# CellRange：表示单元格区域。
# sheet.cell(row, col) 访问即创建，sheet.max_row、sheet.max_column 将随之可能发生变化。
# cell.parent 是 cell 所在 sheet 的引用。


def is_merged_cell(cell: Cell) -> Optional[Union[Cell, MergedCell]]:
    """
    判断给定单元格是否为合并单元格。

    :param cell: 一个单元格。
    :return: 如果它是合并单元格，则返回它本身（MergedCell 或 Cell（如果它是 start_cell）），否则返回 None。
    """
    return cell if isinstance(cell, MergedCell) else None


def is_merged_start_cell(merged_cell: Cell) -> Optional[Cell]:
    """
    判断给定单元格是否为合并单元格的左上角的单元格（即：start_cell）。

    :param merged_cell: 一个合并单元格。
    :return: 如果它是合并单元格的 start_cell，则返回这个 start_cell，否则返回 None。
    """
    for merged in merged_cell.parent.merged_cells:
        if merged.min_row <= merged_cell.row <= merged.max_row \
                and merged.min_col <= merged_cell.column <= merged.max_col:
            return merged.start_cell if merged_cell == merged.start_cell else None
    return None


def get_merged_cell_range(merged_cell: Cell) -> Optional[CellRange]:
    """
    获取合并单元格的区域单元格（即所属的 CellRange）。

    :param merged_cell: 一个合并单元格。
    :return: 如果 merged_cell 不是合并单元格，则返回 None。
    """
    if is_merged_cell(merged_cell):
        for merged in merged_cell.parent.merged_cells:
            if merged.min_row <= merged_cell.row <= merged.max_row \
                    and merged.min_col <= merged_cell.column <= merged.max_col:
                return merged
    return None


def get_cell_by_coord(sheet: Worksheet, row: int, col: int):
    """
    获取工作表中指定行列位置上的单元格。
    
    :param sheet: 指定工作表。
    :param row: 行号。
    :param col: 列号。
    :return: 如果是合并单元格（MergedCell），则返回它的左上角的单元格。
    """
    for merged in sheet.merged_cells:
        if merged.min_row <= row <= merged.max_row and merged.min_col <= col <= merged.max_col:
            return merged.start_cell
    return sheet.cell(row, col)


def get_cell(cell: Cell):
    """
    获取单元格。
    
    :param cell: 一个单元格。
    :return: 如果是合并单元格（MergedCell），则返回它的左上角的单元格，否则返回它本身。
    """
    return get_cell_by_coord(cell.parent, cell.row, cell.column)


def get_cell_value_by_coord(sheet: Worksheet, row: int, col: int):
    """
    获取工作表中指定行列位置上的单元格的内容。

    :param sheet: 指定工作表。
    :param row: 行号。
    :param col: 列号。
    :return: 如果是合并单元格（MergedCell），则返回左上角的单元格的内容。
    """
    return get_cell_by_coord(sheet, row, col).value


def get_cell_value(cell: Cell):
    """
    获取单元格的内容。

    :param cell: 一个单元格。
    :return: 如果是合并单元格（MergedCell），则返回它的左上角的单元格的内容。
    """
    return get_cell_value_by_coord(cell.parent, cell.row, cell.column)


def get_round_cells(cell: Cell, show_merged_cell_value=True):
    """
    获取给定单元格（支持合并单元格）周围（左、上、右、下）的单元格（支持合并单元格）区域。

    :param cell: 一个单元格。
    :param show_merged_cell_value: 是否显示合并单元格。
    :return: 一个元组，四个参数分别表示给定单元格邻近的左、上、右、下的单元格区域。
        如果 show_merged_cell_value=True，并且邻近的单元格如果是合并单元格，则返回它的左上角的单元格（start_cell），否则返回它本身（MergedCell）。
        如果希望以人类视角处理这些单元格，保持 show_merged_cell_value=True，然后消除重复的 start_cell。
    """
    left = ()
    up = ()
    right = ()
    down = ()
    sheet = cell.parent
    merged = is_merged_cell(cell)
    merged_range = get_merged_cell_range(cell)

    if merged:
        if safe_col(cell.column, Direction.LEFT):
            left = (
                get_cell_by_coord(
                    sheet,
                    cell.row,
                    cell.column -
                    1) if show_merged_cell_value else sheet.cell(
                    cell.row,
                    cell.column -
                    1) for (
                    cell.row,
                    cell.column) in merged_range.left)
        if safe_row(cell.row, Direction.UP):
            up = (
                get_cell_by_coord(
                    sheet,
                    cell.row -
                    1,
                    cell.column) if show_merged_cell_value else sheet.cell(
                    cell.row -
                    1,
                    cell.column) for (
                    cell.row,
                    cell.column) in merged_range.top)
        if safe_col(cell.column, Direction.RIGHT):
            right = (
                get_cell_by_coord(
                    sheet,
                    cell.row,
                    cell.column +
                    1) if show_merged_cell_value else sheet.cell(
                    cell.row,
                    cell.column +
                    1) for (
                    cell.row,
                    cell.column) in merged_range.right)
        if safe_row(cell.row, Direction.DOWN):
            down = (
                get_cell_by_coord(
                    sheet,
                    cell.row +
                    1,
                    cell.column) if show_merged_cell_value else sheet.cell(
                    cell.row +
                    1,
                    cell.column) for (
                    cell.row,
                    cell.column) in merged_range.bottom)
    else:
        if safe_col(cell.column, Direction.LEFT):
            left = get_cell_by_coord(
                sheet,
                cell.row,
                cell.column -
                1) if show_merged_cell_value else sheet.cell(
                cell.row,
                cell.column -
                1)
        if safe_row(cell.row, Direction.UP):
            up = get_cell_by_coord(
                sheet,
                cell.row - 1,
                cell.column) if show_merged_cell_value else sheet.cell(
                cell.row - 1,
                cell.column)
        if safe_col(cell.column, Direction.RIGHT):
            right = get_cell_by_coord(
                sheet,
                cell.row,
                cell.column +
                1) if show_merged_cell_value else sheet.cell(
                cell.row,
                cell.column +
                1)
        if safe_row(cell.row, Direction.DOWN):
            down = get_cell_by_coord(
                sheet,
                cell.row + 1,
                cell.column) if show_merged_cell_value else sheet.cell(
                cell.row + 1,
                cell.column)
    return left, up, right, down


def sheet_range_row(sheet: Worksheet, row: int, offset: Union[int, Direction]):
    """
    获取一个行号经过偏移后的行号，如果超出工作表范围，则返回 0。
    """
    if isinstance(offset, int):
        row += offset
    elif isinstance(offset, Direction):
        row += offset.value[0]

    if 1 <= row <= sheet.max_row:
        return row
    return 0


def safe_row(row: int, offset: Union[int, Direction]):
    """
    获取一个行号经过偏移后的行号，如果超出工作表的理论最大范围（1~1048576），则返回 0。
    """
    if isinstance(offset, int):
        row += offset
    elif isinstance(offset, Direction):
        row += offset.value[0]

    if 1 <= row <= 2 << (20 - 1):
        return row
    return 0


def sheet_range_col(sheet: Worksheet, col: int, offset: Union[int, Direction]):
    """
    获取一个列号经过偏移后的列号，如果超出工作表范围，则返回 0。
    """
    if isinstance(offset, int):
        col += offset
    elif isinstance(offset, Direction):
        col += offset.value[1]

    if 1 <= col <= sheet.max_column:
        return col
    return 0


def safe_col(col: int, offset):
    """
    获取一个列号经过偏移后的列号，如果超出工作表的理论最大范围（1~16384），则返回 0。
    """
    if isinstance(offset, int):
        col += offset
    elif isinstance(offset, Direction):
        col += offset.value[1]

    if 1 <= col <= 2 << (14 - 1):
        return col
    return 0


def get_safe_offset_cell_coord(
        cell: Cell, offset: Union[tuple[int, int], Direction]):
    """
    获取一个单元格经过指定偏移后的安全的新坐标，新坐标如果超出工作表理论最大范围，则返回 None。
    """
    if isinstance(offset, Direction):
        offset = offset.value
    sr = safe_row(cell.row, offset[0])
    sc = safe_col(cell.column, offset[1])
    if sr and sc:
        return sr, sc
    return None


def get_safe_offset_cell(
        cell: Cell, offset: Union[tuple[int, int], Direction]):
    """
    获取一个单元格经过指定偏移后的新单元格，新单元格如果超出工作表理论最大范围，则返回 None。
    """
    coord = get_safe_offset_cell_coord(cell, offset)
    if coord:
        return cell.parent.cell(*coord)
    return None


def get_range_by_title_path(sheet: Worksheet,
                            tpath: str,
                            flags: Union[int, re.RegexFlag] = 0,
                            fuzzy=True):
    """
    基于 TitlePath 获取目标区域，支持正则表达式，标志位作用于全局。

    :param sheet: 目标所在工作表。
    :param tpath: 一种用于定位目标区域的路径表达式，语法：“分隔符 路径1 分隔符 路径2 分隔符 路径N...[n]”，如：“/销售部/业绩[1]$DOWN$”，
        它表示 [销售部] 是一个合并单元格（Title），它的下面一行（$DOWN$）包含一个 [业绩] 的合并单元格（Title），
        目标是 [业绩] 单元格（Title）下方的第一列（[1]）的数值，其中目标行/列语法 [n] 只能出现在路径末尾，其上级路径可以是非合并单元格。
        开头第一个分隔符确定了该表达式的分隔符，它通常为 ‘/’，后续需一致，可根据实际情况更换，但不建议使用 ‘$’ 等会产生歧义的符号作为分隔符；
        结尾 $LEFT$、$UP$、$RIGHT$、$DOWN$ 表示路径的寻找方向与目标区域的方向，其中 $DOWN$ 是缺省值，可以省略，方向标识符必须全字母大写。
        综上，一个最基本的 TPath 至少包含开头的分隔符和一个定位 Title，如：“/小计”。
    :param flags: 标志位，用于控制正则表达式的匹配方式，如：是否区分大小写，多行匹配等等。
    :param fuzzy: 查找模式。True：模糊匹配，False：精确匹配。
    :return: 返回目标区域单元格。
    """
    order = parse_tpath_order(tpath)
    d = parse_tpath_direction(tpath)

    tail_cell = find_cell_by_tpath(sheet, tpath, flags, fuzzy)

    if not tail_cell:
        return None

    tail_cell = get_cell_by_coord(sheet, tail_cell.row, tail_cell.column)
    start_cell = get_safe_offset_cell(tail_cell, d)

    if d == Direction.UP or d == Direction.DOWN:
        start_cell = get_safe_offset_cell(
            start_cell, (0, int(Direction.RIGHT.value[1]) + order - 2))
    else:
        start_cell = get_safe_offset_cell(
            start_cell, (int(Direction.DOWN.value[0]) + order - 2, 0))

    yield start_cell

    next_cell = start_cell

    while True:
        next_cell_coord = get_safe_offset_cell_coord(next_cell, d)
        if not next_cell_coord or (
                d == Direction.DOWN and next_cell_coord[0] > sheet.max_row) or (
                d == Direction.RIGHT and next_cell_coord[1] > sheet.max_column):
            return
        next_cell = get_safe_offset_cell(next_cell, d)
        yield next_cell


def parse_tpath_sep(tpath: str):
    """
    解析 TPath 的分隔符。

    :param tpath: 一个 TPath 表达式。
    :return: 如 “/小计”，则返回：/。
    """
    return tpath[0]


def parse_tpath_tpaths(tpath: str):
    """
    解析 TPath 的路径部分。

    :param tpath: 一个 TPath 表达式。
    :return: 如 TPath：“/销售部/业绩[3]$RIGHT$”，则返回：[销售部, 业绩]。
    """
    sep = parse_tpath_sep(tpath)
    order = parse_tpath_order(tpath)
    tpath = tpath[1:]
    tpath = tpath.removesuffix('$LEFT$').removesuffix(
        '$UP$').removesuffix('$RIGHT$').removesuffix('$DOWN$')
    tpath = tpath.removesuffix('[' + str(order) + ']')
    tpaths = tpath.split(sep)
    return tpaths


def parse_tpath_order(tpath: str):
    """
    解析 TPath 的定位序号。

    :param tpath: 一个 TPath 表达式。
    :return: 如 TPath：“/销售部/业绩[3]$RIGHT$”，则返回：3。
    """
    order = 1
    search = re.search(r'\[(\d+)]', tpath)
    if search:
        order = int(search.group(1))
    return order


def parse_tpath_direction(tpath: str):
    """
    解析 TPath 的查找方向（目标区域方向）。

    :param tpath: 一个 TPath 表达式。
    :return: 如 TPath：“/销售部/业绩[3]$RIGHT$”，则返回：Direction.RIGHT。
    """
    d = Direction.DOWN
    if tpath.endswith('$LEFT$'):
        d = Direction.LEFT
    elif tpath.endswith('$UP$'):
        d = Direction.UP
    elif tpath.endswith('$RIGHT$'):
        d = Direction.RIGHT
    elif tpath.endswith('$DOWN$'):
        d = Direction.DOWN
    return d


def parse_tpath(tpath: str):
    """
    解析 TPath。

    :param tpath: 一个 TPath 表达式。
    :return: 一个元组，四个参数分别表示 TPath 的分隔符、路径部分、定位序号、查找方向。
    """
    sep = tpath[0]
    tpath = tpath[1:]
    d = Direction.DOWN

    if tpath.endswith('$LEFT$'):
        d = Direction.LEFT
        tpath = tpath.removesuffix('$LEFT$')
    elif tpath.endswith('$UP$'):
        d = Direction.UP
        tpath = tpath.removesuffix('$UP$')
    elif tpath.endswith('$RIGHT$'):
        d = Direction.RIGHT
        tpath = tpath.removesuffix('$RIGHT$')
    elif tpath.endswith('$DOWN$'):
        d = Direction.DOWN
        tpath = tpath.removesuffix('$DOWN$')

    order = 1
    search = re.search(r'\[(\d+)]$', tpath)
    if search:
        order = int(search.group(1))
        tpath = tpath.removesuffix(search.group(0))

    tpaths = tpath.split(sep)
    return sep, tpaths, order, d


def find_cell_by_tpath(sheet: Worksheet,
                       tpath: str,
                       flags: Union[int, re.RegexFlag] = 0,
                       fuzzy=True) -> Optional[Union[Cell, MergedCell]]:
    """
    基于 TPath 在工作表中查找目标单元格。支持正则表达式，标志位作用于全局。

    :param sheet: 待查找的工作表。
    :param tpath: 一个 TitlePath。
    :param flags: 标志位，用于控制正则表达式的匹配方式，如：是否区分大小写，多行匹配等等。
    :param fuzzy: 查找模式。True：模糊匹配，False：精确匹配。
    :return: 如果目标单元格是合并单元格，则返回 MergedCell。
    """
    tpaths = parse_tpath_tpaths(tpath)
    d = parse_tpath_direction(tpath)
    for rows in sheet.rows:
        for cell in rows:
            next_cell = cell
            is_find = False
            for i in range(len(tpaths)):
                p = tpaths[i]
                next_value = str(
                    get_cell_value_by_coord(
                        sheet,
                        next_cell.row,
                        next_cell.column))
                match = re.search(p, next_value, flags)
                if match:
                    find = match.group(0)
                    if fuzzy:
                        is_find = find in next_value
                    else:
                        is_find = find == next_value
                if i < len(tpaths) - 1:
                    next_cell = get_safe_offset_cell(next_cell, d)
                if not next_cell:
                    is_find = False
                    break
            if is_find:
                return next_cell
    return None


def find_cells_by_regex(
        sheet: Worksheet,
        pattern: str,
        flags: Union[int, re.RegexFlag] = 0,
        fuzzy=True,
        h_v=True,
        show_merged_cell_value=True):
    """
    基于正则表达式在工作表中查找单元格。

    :param sheet: 待查找的工作表。
    :param pattern: 正则表达式。
    :param flags: 标志位，用于控制正则表达式的匹配方式，如：是否区分大小写，多行匹配等等。
    :param fuzzy: 查找模式。True：模糊匹配，False：精确匹配。
    :param h_v: 查找方向。True：逐行查找，False：逐列查找。
    :param show_merged_cell_value: 如果希望以人类视角处理这些单元格，保持 show_merged_cell_value=True，然后消除重复的 start_cell。
    """
    pat = re.compile(pattern, flags)
    for rows_or_cols in (sheet.rows if h_v else sheet.columns):
        for c_r_cell in rows_or_cols:
            value = get_cell_value(c_r_cell) if show_merged_cell_value else sheet.cell(
                c_r_cell.row, c_r_cell.column).value
            if value is None:
                value = ''
            value = str(value)
            match = pat.search(value, flags)
            if match:
                find = match.group(0)
                if fuzzy:
                    if find in value:
                        yield c_r_cell
                else:
                    if find == value:
                        yield c_r_cell


def find_records_by_regex(
        sheet: Worksheet,
        pattern: str,
        flags: Union[int, re.RegexFlag] = 0,
        fuzzy=True,
        show_merged_cell_value=True):
    """
    基于正则表达式在工作表中查找记录（整行）。

    :param sheet: 待查找的工作表。
    :param pattern: 正则表达式。
    :param flags: 标志位，用于控制正则表达式的匹配方式，如：是否区分大小写，多行匹配等等。
    :param fuzzy: 查找模式。True：模糊匹配，False：精确匹配。
    :param show_merged_cell_value: 如果希望以人类视角处理这些单元格，保持 show_merged_cell_value=True。
    :return: 一个元组，(行号, 记录)。
    """
    row_num = 0
    pat = re.compile(pattern, flags)
    for rows in sheet.rows:
        row_num += 1
        for cell in rows:
            value = get_cell_value(cell) if show_merged_cell_value else sheet.cell(
                cell.row, cell.column).value
            if value is None:
                value = ''
            value = str(value)
            match = pat.search(value, flags)
            if match:
                find = match.group(0)
                if fuzzy:
                    if find in value:
                        yield row_num, rows
                        break
                else:
                    if find == value:
                        yield row_num, rows
                        break