book_rag_server/utils/excel_util.py

"""
Excel 工具类

该文件提供 Excel 解析功能，支持：
- 解析 .xlsx 和 .xls 格式的 Excel 文件
- 将 Excel 数据转换为 List[Dict[str, Any]] 格式
- 支持指定工作表
- 支持自定义表头行
"""

from typing import List, Dict, Any, Optional
import os


try:
    from openpyxl import load_workbook
    from openpyxl.worksheet.worksheet import Worksheet
except ImportError:
    raise ImportError("openpyxl 库未安装，请使用 pip install openpyxl 安装")

try:
    import xlrd
    from xlrd.sheet import Sheet
except ImportError:
    raise ImportError("xlrd 库未安装，请使用 pip install xlrd 安装")


class ExcelUtil:
    """
    Excel 工具类，用于解析 Excel 文件
    """
    
    @staticmethod
    def _parse_xlsx(file_path: str, sheet_name: Optional[str] = None, header_row: int = 0) -> List[Dict[str, Any]]:
        """
        解析 .xlsx 格式的 Excel 文件
        
        Args:
            file_path: Excel 文件路径
            sheet_name: 工作表名称，None 表示使用第一个工作表
            header_row: 表头所在行索引，默认第 0 行
            
        Returns:
            List[Dict[str, Any]]: 解析后的 Excel 数据
        """
        # 加载工作簿
        workbook = load_workbook(filename=file_path, read_only=True)
        
        # 获取工作表
        if sheet_name:
            sheet: Worksheet = workbook[sheet_name]
        else:
            sheet: Worksheet = workbook.active
        
        # 获取所有行数据
        rows = list(sheet.iter_rows(values_only=True))
        
        # 检查数据是否为空
        if not rows:
            return []
        
        # 检查表头行索引是否合法
        if header_row >= len(rows):
            raise ValueError(f"表头行索引 {header_row} 超出数据总行数 {len(rows)}")
        
        # 获取表头
        headers = rows[header_row]
        
        # 检查表头是否为空
        if not headers:
            return []
        
        # 解析数据行
        data: List[Dict[str, Any]] = []
        for row in rows[header_row + 1:]:
            if not row or all(cell is None for cell in row):
                continue
            
            # 创建数据字典
            row_data: Dict[str, Any] = {}
            for i, header in enumerate(headers):
                if i < len(row):
                    row_data[header] = row[i]
                else:
                    row_data[header] = None
            
            data.append(row_data)
        
        workbook.close()
        return data
    
    @staticmethod
    def _parse_xls(file_path: str, sheet_name: Optional[str] = None, header_row: int = 0) -> List[Dict[str, Any]]:
        """
        解析 .xls 格式的 Excel 文件
        
        Args:
            file_path: Excel 文件路径
            sheet_name: 工作表名称，None 表示使用第一个工作表
            header_row: 表头所在行索引，默认第 0 行
            
        Returns:
            List[Dict[str, Any]]: 解析后的 Excel 数据
        """
        # 打开工作簿
        workbook = xlrd.open_workbook(file_path)
        
        # 获取工作表
        if sheet_name:
            sheet: Sheet = workbook.sheet_by_name(sheet_name)
        else:
            sheet: Sheet = workbook.sheet_by_index(0)
        
        # 获取所有行数据
        rows = []
        for i in range(sheet.nrows):
            rows.append(sheet.row_values(i))
        
        # 检查数据是否为空
        if not rows:
            return []
        
        # 检查表头行索引是否合法
        if header_row >= len(rows):
            raise ValueError(f"表头行索引 {header_row} 超出数据总行数 {len(rows)}")
        
        # 获取表头
        headers = rows[header_row]
        
        # 检查表头是否为空
        if not headers:
            return []
        
        # 解析数据行
        data: List[Dict[str, Any]] = []
        for row in rows[header_row + 1:]:
            if not row or all(cell == '' or cell is None for cell in row):
                continue
            
            # 创建数据字典
            row_data: Dict[str, Any] = {}
            for i, header in enumerate(headers):
                if i < len(row):
                    row_data[header] = row[i]
                else:
                    row_data[header] = None
            
            data.append(row_data)
        
        return data
    
    @staticmethod
    def parse_excel(file_path: str, sheet_name: Optional[str] = None, header_row: int = 0) -> List[Dict[str, Any]]:
        """
        解析 Excel 文件，支持 .xlsx 和 .xls 格式
        
        Args:
            file_path: Excel 文件路径
            sheet_name: 工作表名称，None 表示使用第一个工作表
            header_row: 表头所在行索引，默认第 0 行
            
        Returns:
            List[Dict[str, Any]]: 解析后的 Excel 数据
        
        Raises:
            ValueError: 不支持的文件格式
            FileNotFoundError: 文件不存在
        """
        # 检查文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件 {file_path} 不存在")
        
        # 获取文件扩展名
        file_ext = os.path.splitext(file_path)[1].lower()
        
        # 根据文件格式选择解析方法
        if file_ext == '.xlsx':
            return ExcelUtil._parse_xlsx(file_path, sheet_name, header_row)
        elif file_ext == '.xls':
            return ExcelUtil._parse_xls(file_path, sheet_name, header_row)
        else:
            raise ValueError(f"不支持的文件格式 {file_ext}，只支持 .xlsx 和 .xls 格式")
    
    @staticmethod
    def get_sheet_names(file_path: str) -> List[str]:
        """
        获取 Excel 文件中的所有工作表名称
        
        Args:
            file_path: Excel 文件路径
            
        Returns:
            List[str]: 工作表名称列表
        
        Raises:
            ValueError: 不支持的文件格式
            FileNotFoundError: 文件不存在
        """
        # 检查文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件 {file_path} 不存在")
        
        # 获取文件扩展名
        file_ext = os.path.splitext(file_path)[1].lower()
        
        # 根据文件格式选择方法
        if file_ext == '.xlsx':
            workbook = load_workbook(filename=file_path, read_only=True)
            sheet_names = workbook.sheetnames
            workbook.close()
            return sheet_names
        elif file_ext == '.xls':
            workbook = xlrd.open_workbook(file_path)
            sheet_names = workbook.sheet_names()
            return sheet_names
        else:
            raise ValueError(f"不支持的文件格式 {file_ext}，只支持 .xlsx 和 .xls 格式")


# 单例模式实例
excel_util = ExcelUtil()