excel_tool/function_calling/excel_processor.py

import pandas as pd
from typing import List
import logging
from mylib.logging_config import setup_logging

# Setup custom logging
setup_logging()
logger = logging.getLogger('excel_tool')

class ExcelProcessor:
    def __init__(self):
        """Initialize Excel processor"""
        pass

    def _column_letter_to_index(self, col: str) -> int:
        """将Excel列字母转换为索引（A=0, B=1,...）"""
        index = 0
        for char in col.upper():
            if not 'A' <= char <= 'Z':
                raise ValueError(f"无效的列字母: {col}")
            index = index * 26 + (ord(char) - ord('A') + 1)
        return index - 1

    def read_excel_file(self, file_path: str) -> pd.DataFrame:
        """读取Excel文件并返回DataFrame"""
        try:
            df = pd.read_excel(file_path)
            logger.info(f"成功读取文件: {file_path}")
            return df
        except Exception as e:
            logger.error(f"读取文件失败: {file_path}, 错误: {str(e)}")
            raise

    def save_excel_file(self, df: pd.DataFrame, output_path: str) -> None:
        """将DataFrame保存为Excel文件"""
        try:
            df.to_excel(output_path, index=False)
            logger.info(f"成功保存文件: {output_path}")
        except Exception as e:
            logger.error(f"保存文件失败: {output_path}, 错误: {str(e)}")
            raise

    def insert_column(self, df: pd.DataFrame, ref_column: str, new_column_name: str, position: str = 'right') -> pd.DataFrame:
        """
        在指定列旁边插入空列
        
        Args:
            df (pd.DataFrame): 输入的DataFrame
            ref_column (str): 参考列名或Excel列字母（如'A', 'B'）
            new_column_name (str): 新列名
            position (str): 插入位置，'left'或'right'，默认为'right'
            
        Returns:
            pd.DataFrame: 包含新列的DataFrame
        """
        try:
            # 如果ref_column是字母，转换为列索引
            if ref_column.isalpha():
                col_index = self._column_letter_to_index(ref_column)
                ref_column = df.columns[col_index]
            
            # 获取参考列的位置
            ref_index = df.columns.get_loc(ref_column)
            
            # 计算插入位置
            insert_index = ref_index + 1 if position == 'right' else ref_index
            
            # 插入新列
            df.insert(insert_index, new_column_name, '')
            
            logger.info(f"成功在列 {ref_column} 的{position}插入新列 {new_column_name}")
            return df
            
        except Exception as e:
            logger.error(f"插入列失败: {str(e)}")
            raise

if __name__ == '__main__':
    # 测试代码
    processor = ExcelProcessor()
    df = processor.read_excel_file('/home/mrh/code/excel_tool/temp/测试.xlsx')
    df = processor.insert_column(df, 'B', 'New Column')
    processor.save_excel_file(df, '/home/mrh/code/excel_tool/temp/测试_process.xlsx')