Prechádzať zdrojové kódy

feat: add header row parameter to CSV reading functions

mrh (aider) 1 rok pred
rodič
commit
a42b05cb23
1 zmenil súbory, kde vykonal 31 pridanie a 4 odobranie
  1. 31 4
      mylib/translate_utils.py

+ 31 - 4
mylib/translate_utils.py

@@ -163,20 +163,47 @@ def process_batch_translations(df: pd.DataFrame,
         logger.error(f"批量翻译时出错: {e}")
         raise
 
+def read_csv_with_header(file_path: str, header_row: int = 1, encoding: str = None) -> pd.DataFrame:
+    """读取CSV文件并正确处理标题行
+    
+    Args:
+        file_path: CSV文件路径
+        header_row: 标题行号(从0开始),默认为1(第2行)
+        encoding: 文件编码
+        
+    Returns:
+        pandas DataFrame
+    """
+    try:
+        # 读取所有数据
+        data = read_csv(file_path, encoding)
+        
+        # 确保header_row在有效范围内
+        if header_row >= len(data):
+            raise ValueError(f"标题行 {header_row} 超出文件范围")
+            
+        # 使用指定行作为列名,前面的行丢弃
+        df = pd.DataFrame(data[header_row+1:], columns=data[header_row])
+        
+        logger.info(f"成功读取CSV文件,使用第{header_row+1}行作为标题行")
+        return df
+    except Exception as e:
+        logger.error(f"读取CSV文件时出错: {e}")
+        raise
+
 def main():
     output_dir = Path('temp')
     input_file = output_dir/"测试.csv"
     output_file = output_dir/"processed_测试.csv"
     
     # 使用自定义编码检测读取CSV文件
-    data = read_csv(input_file)
-    df = pd.DataFrame(data[1:], columns=data[0])
+    df = read_csv_with_header(input_file, header_row=1)  # 使用第2行作为标题行
     
     # 提取列数据
-    extract_column_data(df, 'B', start_row=2, header_row=2)  # 示例:从第3行开始提取第2列(即'B'列)的数据
+    extract_column_data(df, 'B', start_row=2, header_row=1)  # 示例:从第3行开始提取第2列(即'B'列)的数据
     
     # 插入空列
-    df = insert_empty_columns(df, ['B'], header_row=2)  # 示例:在'B'列后插入空列
+    df = insert_empty_columns(df, ['B'], header_row=1)  # 示例:在'B'列后插入空列
     
     # 处理翻译
     # df, _ = process_batch_translations(df, '搜索词')