301 lines
9.8 KiB
Python
301 lines
9.8 KiB
Python
import csv
|
||
import os
|
||
import time
|
||
from typing import List, Dict, Any, Optional
|
||
|
||
|
||
class CSVTool:
|
||
def __init__(self, csv_file_name: str, headers: List[str]):
|
||
"""
|
||
初始化CSV工具
|
||
|
||
Args:
|
||
csv_file_name (str): CSV文件名
|
||
headers (List[str]): 表头列表
|
||
"""
|
||
self.csv_file_name = csv_file_name
|
||
self.headers = headers
|
||
|
||
def init_csv_file(self):
|
||
"""
|
||
初始化CSV文件，如果文件不存在则创建并写入表头
|
||
"""
|
||
if not os.path.exists(self.csv_file_name):
|
||
with open(self.csv_file_name, 'w', encoding='utf-8', newline='') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow(self.headers)
|
||
|
||
# 文件存在，检查是否为空或只有空行
|
||
try:
|
||
with open(self.csv_file_name, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
|
||
# 如果文件内容为空或只包含空白字符（空格、换行符等）
|
||
if not content.strip():
|
||
with open(self.csv_file_name, 'w', encoding='utf-8', newline='') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow(self.headers)
|
||
return
|
||
|
||
# 文件有内容，检查第一行是否为表头
|
||
with open(self.csv_file_name, 'r', encoding='utf-8') as f:
|
||
reader = csv.reader(f)
|
||
first_row = next(reader, None)
|
||
|
||
# 如果第一行不是预期的表头，则重新写入表头和原有内容
|
||
if first_row != self.headers:
|
||
lines = content.strip().split('\n')
|
||
# 过滤掉空行和只包含空白字符的行
|
||
non_empty_lines = [line for line in lines if line.strip()]
|
||
|
||
with open(self.csv_file_name, 'w', encoding='utf-8', newline='') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow(self.headers)
|
||
# 如果有非空内容，则写入
|
||
if non_empty_lines:
|
||
f.write('\n'.join(non_empty_lines) + '\n')
|
||
|
||
except Exception as e:
|
||
print(f"检查/更新表头时出错: {e}")
|
||
|
||
|
||
def get_existing_data(self, unique_titles: List[str]) -> set:
|
||
"""
|
||
读取现有数据，用于去重检查
|
||
|
||
Args:
|
||
unique_titles (List[str]): 用于唯一性检查的列名列表
|
||
|
||
Returns:
|
||
set: 包含唯一标识符元组的集合
|
||
"""
|
||
existing_data = set()
|
||
|
||
if not os.path.exists(self.csv_file_name):
|
||
return existing_data
|
||
|
||
try:
|
||
with open(self.csv_file_name, 'r', encoding='utf-8') as f:
|
||
reader = csv.reader(f)
|
||
header_row = next(reader, None) # 读取表头
|
||
|
||
if header_row is None:
|
||
return existing_data
|
||
|
||
# 获取唯一列的索引
|
||
unique_indices = []
|
||
for title in unique_titles:
|
||
try:
|
||
index = header_row.index(title)
|
||
unique_indices.append(index)
|
||
except ValueError:
|
||
print(f"警告: 表头中未找到列 '{title}'")
|
||
continue
|
||
|
||
# 读取数据行
|
||
for row in reader:
|
||
if len(row) >= len(header_row): # 确保行数据完整
|
||
# 提取唯一标识符
|
||
unique_values = tuple(row[i] if i < len(row) else "" for i in unique_indices)
|
||
existing_data.add(unique_values)
|
||
|
||
except Exception as e:
|
||
print(f"读取CSV文件时出错: {e}")
|
||
|
||
return existing_data
|
||
|
||
def save_data(self, data_list: List[Dict[str, Any]], unique_titles: List[str], create_time: bool = True) -> int:
|
||
"""
|
||
将数据保存到CSV文件中，自动去重
|
||
|
||
Args:
|
||
data_list (List[Dict[str, Any]]): 要保存的数据列表
|
||
unique_titles (List[str]): 用于唯一性检查的列名列表
|
||
create_time (bool): 是否自动添加创建时间，默认为True
|
||
|
||
Returns:
|
||
int: 实际写入的行数
|
||
"""
|
||
if not data_list:
|
||
print('数据列表为空，没有数据可写入')
|
||
return 0
|
||
|
||
# 初始化文件
|
||
self.init_csv_file()
|
||
|
||
# 获取现有数据用于去重
|
||
existing_data = self.get_existing_data(unique_titles)
|
||
|
||
# 准备写入的数据
|
||
rows_to_write = []
|
||
written_count = 0
|
||
|
||
for data_node in data_list:
|
||
# 构建唯一标识符元组
|
||
unique_values = tuple(data_node.get(title, "") for title in unique_titles)
|
||
|
||
# 检查是否已存在
|
||
if unique_values in existing_data:
|
||
continue # 跳过已存在的数据
|
||
|
||
# 构建行数据
|
||
row_data = []
|
||
for header in self.headers:
|
||
if header == 'create_time' and create_time:
|
||
# 自动添加创建时间
|
||
row_data.append(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
|
||
else:
|
||
# 从数据节点获取对应值
|
||
row_data.append(data_node.get(header, ""))
|
||
|
||
rows_to_write.append(row_data)
|
||
existing_data.add(unique_values) # 添加到已存在数据集合中，避免本次写入中的重复
|
||
written_count += 1
|
||
|
||
# 写入数据
|
||
if rows_to_write:
|
||
with open(self.csv_file_name, 'a', encoding='utf-8', newline='') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerows(rows_to_write)
|
||
|
||
print(f"成功写入 {written_count} 行数据到 {self.csv_file_name}")
|
||
else:
|
||
print("没有新数据需要写入")
|
||
|
||
return written_count
|
||
|
||
def query_data(self, filter_func=None) -> List[Dict[str, str]]:
|
||
"""
|
||
查询CSV文件中的数据
|
||
|
||
Args:
|
||
filter_func (callable, optional): 过滤函数，接受一行数据字典作为参数，返回True/False
|
||
|
||
Returns:
|
||
List[Dict[str, str]]: 查询结果列表
|
||
"""
|
||
if not os.path.exists(self.csv_file_name):
|
||
return []
|
||
|
||
result = []
|
||
|
||
try:
|
||
with open(self.csv_file_name, 'r', encoding='utf-8') as f:
|
||
reader = csv.reader(f)
|
||
header_row = next(reader, None) # 读取表头
|
||
|
||
if header_row is None:
|
||
return result
|
||
|
||
for row in reader:
|
||
if len(row) >= len(header_row): # 确保行数据完整
|
||
# 将行数据转换为字典
|
||
row_dict = dict(zip(header_row, row))
|
||
|
||
# 应用过滤条件
|
||
if filter_func is None or filter_func(row_dict):
|
||
result.append(row_dict)
|
||
|
||
except Exception as e:
|
||
print(f"查询CSV文件时出错: {e}")
|
||
|
||
return result
|
||
|
||
def query_by_conditions(self, **kwargs) -> List[Dict[str, str]]:
|
||
"""
|
||
根据条件查询数据
|
||
|
||
Args:
|
||
**kwargs: 查询条件，键值对形式
|
||
|
||
Returns:
|
||
List[Dict[str, str]]: 查询结果列表
|
||
"""
|
||
|
||
def filter_func(row_dict):
|
||
for key, value in kwargs.items():
|
||
if key in row_dict and row_dict[key] != value:
|
||
return False
|
||
return True
|
||
|
||
return self.query_data(filter_func)
|
||
|
||
def get_all_data(self) -> List[Dict[str, str]]:
|
||
"""
|
||
获取所有数据
|
||
|
||
Returns:
|
||
List[Dict[str, str]]: 所有数据列表
|
||
"""
|
||
return self.query_data()
|
||
|
||
|
||
# 保持向后兼容的函数
|
||
def save_to_csv(filter_list: List[Dict[str, Any]],
|
||
csv_file_name: str = 'company_search_result_data.csv',
|
||
headers: List[str] = None,
|
||
unique_titles: List[str] = None) -> bool:
|
||
"""
|
||
将结果追加写入csv文件中（向后兼容函数）
|
||
|
||
Args:
|
||
filter_list: 需要写入的数据列表
|
||
csv_file_name: CSV文件名
|
||
headers: 表头列表
|
||
unique_titles: 用于唯一性检查的列名列表
|
||
|
||
Returns:
|
||
bool: 是否成功写入
|
||
"""
|
||
if headers is None:
|
||
headers = ['title', 'url', 'web_site_type', 'request_url', 'company_name', 'create_time']
|
||
|
||
if unique_titles is None:
|
||
unique_titles = ['company_name', 'web_site_type']
|
||
|
||
try:
|
||
csv_tool = CSVTool(csv_file_name, headers)
|
||
written_count = csv_tool.save_data(filter_list, unique_titles)
|
||
return written_count > 0
|
||
except Exception as e:
|
||
print(f"保存CSV时出错: {e}")
|
||
return False
|
||
|
||
|
||
# 使用示例:
|
||
if __name__ == "__main__":
|
||
# 示例数据
|
||
sample_data = [
|
||
{
|
||
'title': '测试公司1',
|
||
'url': 'https://example.com/1',
|
||
'web_site_type': 'aiqicha',
|
||
'request_url': 'https://bing.com/search?q=测试公司1',
|
||
'company_name': '测试公司1'
|
||
},
|
||
{
|
||
'title': '测试公司2',
|
||
'url': 'https://example.com/2',
|
||
'web_site_type': 'qcc',
|
||
'request_url': 'https://bing.com/search?q=测试公司2',
|
||
'company_name': '测试公司2'
|
||
}
|
||
]
|
||
|
||
# 创建CSV工具实例
|
||
csv_tool = CSVTool(
|
||
csv_file_name='test_data.csv',
|
||
headers=['title', 'url', 'web_site_type', 'request_url', 'company_name', 'create_time']
|
||
)
|
||
|
||
# 保存数据
|
||
csv_tool.save_data(sample_data, unique_titles=['company_name', 'web_site_type'])
|
||
|
||
# 查询所有数据
|
||
all_data = csv_tool.get_all_data()
|
||
print("所有数据:", all_data)
|
||
|
||
# 根据条件查询
|
||
filtered_data = csv_tool.query_by_conditions(web_site_type='aiqicha')
|
||
print("查询结果:", filtered_data)
|