316 lines
10 KiB
Python
316 lines
10 KiB
Python
"""
|
||
CSV文件处理工具
|
||
提供CSV文件的读取、写入、验证和备份功能
|
||
"""
|
||
|
||
import os
|
||
import csv
|
||
import uuid
|
||
import json
|
||
import shutil
|
||
import datetime
|
||
from typing import List, Dict, Any, Optional, Callable
|
||
|
||
|
||
class CSVHandler:
|
||
"""
|
||
CSV文件处理类，提供读写、验证和备份功能
|
||
"""
|
||
|
||
def __init__(self, file_path: str, backup_dir: str = None, required_fields: List[str] = None):
|
||
"""
|
||
初始化CSV处理器
|
||
|
||
Args:
|
||
file_path: CSV文件路径
|
||
backup_dir: 备份目录路径，默认为文件所在目录下的backups文件夹
|
||
required_fields: 必填字段列表
|
||
"""
|
||
self.file_path = file_path
|
||
self.required_fields = required_fields or []
|
||
|
||
# 设置备份目录
|
||
if backup_dir is None:
|
||
file_dir = os.path.dirname(os.path.abspath(file_path))
|
||
self.backup_dir = os.path.join(file_dir, 'backups')
|
||
else:
|
||
self.backup_dir = backup_dir
|
||
|
||
# 确保备份目录存在
|
||
os.makedirs(self.backup_dir, exist_ok=True)
|
||
|
||
# 确保文件目录存在
|
||
os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True)
|
||
|
||
def file_exists(self) -> bool:
|
||
"""
|
||
检查CSV文件是否存在
|
||
|
||
Returns:
|
||
文件是否存在
|
||
"""
|
||
return os.path.exists(self.file_path)
|
||
|
||
def create_file(self, headers: List[str]) -> None:
|
||
"""
|
||
创建CSV文件并写入表头
|
||
|
||
Args:
|
||
headers: 表头字段列表
|
||
"""
|
||
# 确保所有必填字段都在表头中
|
||
for field in self.required_fields:
|
||
if field not in headers:
|
||
raise ValueError(f"必填字段 '{field}' 不在表头中")
|
||
|
||
# 创建CSV文件并写入表头
|
||
with open(self.file_path, 'w', newline='', encoding='utf-8') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow(headers)
|
||
|
||
def read_data(self) -> List[Dict[str, Any]]:
|
||
"""
|
||
读取CSV文件数据
|
||
|
||
Returns:
|
||
字典列表，每个字典代表一行数据
|
||
"""
|
||
if not self.file_exists():
|
||
return []
|
||
|
||
data = []
|
||
with open(self.file_path, 'r', encoding='utf-8') as f:
|
||
reader = csv.DictReader(f)
|
||
for row in reader:
|
||
# 处理可能的空字符串
|
||
processed_row = {}
|
||
for key, value in row.items():
|
||
if value == '':
|
||
processed_row[key] = None
|
||
else:
|
||
# 尝试解析JSON字段
|
||
if key in ['metrics'] and value:
|
||
try:
|
||
processed_row[key] = json.loads(value)
|
||
except json.JSONDecodeError:
|
||
processed_row[key] = value
|
||
else:
|
||
processed_row[key] = value
|
||
data.append(processed_row)
|
||
return data
|
||
|
||
def write_data(self, data: List[Dict[str, Any]], validate: bool = True) -> None:
|
||
"""
|
||
写入数据到CSV文件
|
||
|
||
Args:
|
||
data: 要写入的数据，字典列表
|
||
validate: 是否验证数据
|
||
"""
|
||
if validate:
|
||
self._validate_data(data)
|
||
|
||
# 在写入前创建备份
|
||
self._create_backup()
|
||
|
||
# 准备写入，如果文件不存在则创建
|
||
if not data:
|
||
return
|
||
|
||
# 获取所有字段（表头）
|
||
all_fields = set()
|
||
for item in data:
|
||
all_fields.update(item.keys())
|
||
|
||
# 确保必填字段在表头中
|
||
for field in self.required_fields:
|
||
if field not in all_fields:
|
||
raise ValueError(f"必填字段 '{field}' 缺失")
|
||
|
||
headers = list(all_fields)
|
||
|
||
# 写入数据
|
||
with open(self.file_path, 'w', newline='', encoding='utf-8') as f:
|
||
writer = csv.DictWriter(f, fieldnames=headers)
|
||
writer.writeheader()
|
||
|
||
for row in data:
|
||
# 处理特殊字段
|
||
processed_row = {}
|
||
for key, value in row.items():
|
||
if value is None:
|
||
processed_row[key] = ''
|
||
elif key in ['metrics'] and isinstance(value, (dict, list)):
|
||
processed_row[key] = json.dumps(value, ensure_ascii=False)
|
||
else:
|
||
processed_row[key] = value
|
||
writer.writerow(processed_row)
|
||
|
||
def _validate_data(self, data: List[Dict[str, Any]]) -> None:
|
||
"""
|
||
验证数据是否符合要求
|
||
|
||
Args:
|
||
data: 要验证的数据
|
||
|
||
Raises:
|
||
ValueError: 当数据验证失败时
|
||
"""
|
||
for i, item in enumerate(data):
|
||
# 检查必填字段
|
||
for field in self.required_fields:
|
||
if field not in item:
|
||
raise ValueError(f"第 {i+1} 行数据缺少必填字段 '{field}'")
|
||
if item[field] is None or item[field] == '':
|
||
raise ValueError(f"第 {i+1} 行数据的必填字段 '{field}' 为空")
|
||
|
||
def _create_backup(self) -> None:
|
||
"""
|
||
创建CSV文件的备份
|
||
"""
|
||
if not self.file_exists():
|
||
return
|
||
|
||
# 创建带时间戳的备份文件名
|
||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
filename = os.path.basename(self.file_path)
|
||
name, ext = os.path.splitext(filename)
|
||
backup_filename = f"{name}_{timestamp}{ext}"
|
||
backup_path = os.path.join(self.backup_dir, backup_filename)
|
||
|
||
# 复制文件
|
||
shutil.copy2(self.file_path, backup_path)
|
||
|
||
def validate_file(self) -> bool:
|
||
"""
|
||
验证CSV文件格式和数据是否正确
|
||
|
||
Returns:
|
||
验证是否通过
|
||
"""
|
||
if not self.file_exists():
|
||
return False
|
||
|
||
try:
|
||
data = self.read_data()
|
||
self._validate_data(data)
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def append_row(self, row: Dict[str, Any], validate: bool = True) -> None:
|
||
"""
|
||
追加一行数据到CSV文件
|
||
|
||
Args:
|
||
row: 要追加的数据行
|
||
validate: 是否验证数据
|
||
"""
|
||
data = self.read_data()
|
||
data.append(row)
|
||
self.write_data(data, validate=validate)
|
||
|
||
def update_row(self, key_field: str, key_value: str,
|
||
new_data: Dict[str, Any], validate: bool = True) -> bool:
|
||
"""
|
||
更新CSV文件中的一行数据
|
||
|
||
Args:
|
||
key_field: 用于识别行的键字段名
|
||
key_value: 要更新的行的键字段值
|
||
new_data: 新的数据，将合并到现有数据
|
||
validate: 是否验证数据
|
||
|
||
Returns:
|
||
是否成功更新了数据
|
||
"""
|
||
data = self.read_data()
|
||
updated = False
|
||
|
||
for i, row in enumerate(data):
|
||
if row.get(key_field) == key_value:
|
||
data[i] = {**row, **new_data}
|
||
updated = True
|
||
break
|
||
|
||
if updated:
|
||
self.write_data(data, validate=validate)
|
||
|
||
return updated
|
||
|
||
def delete_row(self, key_field: str, key_value: str, validate: bool = True) -> bool:
|
||
"""
|
||
删除CSV文件中的一行数据
|
||
|
||
Args:
|
||
key_field: 用于识别行的键字段名
|
||
key_value: 要删除的行的键字段值
|
||
validate: 是否验证数据
|
||
|
||
Returns:
|
||
是否成功删除了数据
|
||
"""
|
||
data = self.read_data()
|
||
original_len = len(data)
|
||
|
||
data = [row for row in data if row.get(key_field) != key_value]
|
||
|
||
if len(data) < original_len:
|
||
self.write_data(data, validate=validate)
|
||
return True
|
||
else:
|
||
return False
|
||
|
||
def find_rows(self, filter_func: Callable[[Dict[str, Any]], bool]) -> List[Dict[str, Any]]:
|
||
"""
|
||
使用过滤函数查找满足条件的行
|
||
|
||
Args:
|
||
filter_func: 过滤函数，接收一个字典作为参数，返回布尔值
|
||
|
||
Returns:
|
||
满足条件的行列表
|
||
"""
|
||
data = self.read_data()
|
||
return [row for row in data if filter_func(row)]
|
||
|
||
def get_row(self, key_field: str, key_value: str) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
获取指定键值的行
|
||
|
||
Args:
|
||
key_field: 键字段名
|
||
key_value: 键字段值
|
||
|
||
Returns:
|
||
找到的行数据，未找到则返回None
|
||
"""
|
||
data = self.read_data()
|
||
for row in data:
|
||
if row.get(key_field) == key_value:
|
||
return row
|
||
return None
|
||
|
||
|
||
def create_default_goal_csv(file_path: str = 'data/goals.csv') -> None:
|
||
"""
|
||
创建默认的目标CSV文件
|
||
|
||
Args:
|
||
file_path: CSV文件路径
|
||
"""
|
||
required_fields = ['id', 'title', 'type', 'status', 'created_time',
|
||
'modified_time', 'keywords', 'progress']
|
||
|
||
headers = ['id', 'title', 'description', 'type', 'status',
|
||
'created_time', 'modified_time', 'start_date', 'due_date',
|
||
'frequency', 'keywords', 'progress', 'related_projects', 'metrics']
|
||
|
||
handler = CSVHandler(file_path, required_fields=required_fields)
|
||
|
||
if not handler.file_exists():
|
||
handler.create_file(headers)
|
||
print(f"已创建默认目标CSV文件: {file_path}")
|
||
else:
|
||
print(f"目标CSV文件已存在: {file_path}")
|