Просмотр исходного кода

feat: Add CSV support and specify Excel engine in file processor

mrh (aider) 1 год назад
Родитель
Сommit
a0e039c714
1 измененных файлов с 31 добавлено и 10 удалено
  1. 31 10
      function_calling/excel_processor.py

+ 31 - 10
function_calling/excel_processor.py

@@ -21,21 +21,37 @@ class ExcelProcessor:
             index = index * 26 + (ord(char) - ord('A') + 1)
         return index - 1
 
-    def read_excel_file(self, file_path: str) -> pd.DataFrame:
-        """读取Excel文件并返回DataFrame"""
+    def read_file(self, file_path: str) -> pd.DataFrame:
+        """
+        读取文件并返回DataFrame
+        支持Excel和CSV文件
+        """
         try:
-            df = pd.read_excel(file_path)
-            logger.info(f"成功读取文件: {file_path}")
+            if file_path.endswith('.csv'):
+                df = pd.read_csv(file_path)
+                logger.info(f"成功读取CSV文件: {file_path}")
+            else:
+                # 对于Excel文件,指定engine参数
+                df = pd.read_excel(file_path, engine='openpyxl')
+                logger.info(f"成功读取Excel文件: {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文件"""
+    def save_file(self, df: pd.DataFrame, output_path: str) -> None:
+        """
+        将DataFrame保存为文件
+        支持Excel和CSV格式
+        """
         try:
-            df.to_excel(output_path, index=False)
-            logger.info(f"成功保存文件: {output_path}")
+            if output_path.endswith('.csv'):
+                df.to_csv(output_path, index=False)
+                logger.info(f"成功保存CSV文件: {output_path}")
+            else:
+                # 对于Excel文件,指定engine参数
+                df.to_excel(output_path, index=False, engine='openpyxl')
+                logger.info(f"成功保存Excel文件: {output_path}")
         except Exception as e:
             logger.error(f"保存文件失败: {output_path}, 错误: {str(e)}")
             raise
@@ -78,6 +94,11 @@ class ExcelProcessor:
 if __name__ == '__main__':
     # 测试代码
     processor = ExcelProcessor()
-    df = processor.read_excel_file('/home/mrh/code/excel_tool/temp/测试.csv')
+    # 测试CSV文件
+    df = processor.read_file('/home/mrh/code/excel_tool/temp/测试.csv')
+    df = processor.insert_column(df, 'B', 'New Column')
+    processor.save_file(df, '/home/mrh/code/excel_tool/temp/测试_process.csv')
+    # 测试Excel文件
+    df = processor.read_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')
+    processor.save_file(df, '/home/mrh/code/excel_tool/temp/测试_process.xlsx')