346 lines
10 KiB
Python
346 lines
10 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
每日热榜 Skill - 数据存储模块
|
||
|
||
功能：
|
||
- 自动保存每日热榜数据
|
||
- 查询历史热榜记录
|
||
- 管理数据文件
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
from datetime import datetime, timedelta
|
||
from typing import Dict, List, Optional, Any
|
||
from config import config
|
||
|
||
|
||
class DataStorage:
|
||
"""数据存储类"""
|
||
|
||
def __init__(self):
|
||
self.data_dir = config.data_dir
|
||
self.auto_save = config.auto_save
|
||
os.makedirs(self.data_dir, exist_ok=True)
|
||
|
||
def save_hot_list(self, source_id: str, data: Dict[str, Any]) -> bool:
|
||
"""
|
||
保存热榜数据到文件
|
||
|
||
Args:
|
||
source_id: 热榜源 ID（如 weibo, zhihu）
|
||
data: 热榜数据字典
|
||
|
||
Returns:
|
||
是否保存成功
|
||
"""
|
||
if not self.auto_save:
|
||
return False
|
||
|
||
try:
|
||
# 获取今日日期的文件路径
|
||
file_path = config.get_data_path(source_id)
|
||
|
||
# 构建存储结构
|
||
storage_data = {
|
||
"source_id": source_id,
|
||
"save_time": datetime.now().isoformat(),
|
||
"update_time": data.get("update_time", ""),
|
||
"total": data.get("total", 0),
|
||
"data": data.get("data", [])
|
||
}
|
||
|
||
# 读取现有数据（如果有）
|
||
existing_data = []
|
||
if os.path.exists(file_path):
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
try:
|
||
existing_data = json.load(f)
|
||
if not isinstance(existing_data, list):
|
||
existing_data = []
|
||
except:
|
||
existing_data = []
|
||
|
||
# 添加新数据到列表
|
||
existing_data.append(storage_data)
|
||
|
||
# 保存到文件
|
||
with open(file_path, 'w', encoding='utf-8') as f:
|
||
json.dump(existing_data, f, ensure_ascii=False, indent=2)
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"[DataStorage] 保存失败: {e}")
|
||
return False
|
||
|
||
def load_history(self, source_id: str = "", days: int = 7) -> Dict[str, List[Dict]]:
|
||
"""
|
||
加载历史热榜数据
|
||
|
||
Args:
|
||
source_id: 热榜源 ID（空则加载所有源）
|
||
days: 加载最近几天的数据
|
||
|
||
Returns:
|
||
按日期组织的热榜数据
|
||
"""
|
||
result = {}
|
||
|
||
if source_id:
|
||
# 加载指定源的历史数据
|
||
source_dir = os.path.join(self.data_dir, source_id)
|
||
if not os.path.exists(source_dir):
|
||
return {}
|
||
|
||
for i in range(days):
|
||
date_str = (datetime.now() - timedelta(days=i)).strftime("%Y-%m-%d")
|
||
file_path = os.path.join(source_dir, f"{date_str}.json")
|
||
|
||
if os.path.exists(file_path):
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
try:
|
||
data = json.load(f)
|
||
result[date_str] = data
|
||
except:
|
||
continue
|
||
|
||
else:
|
||
# 加载所有源的历史数据
|
||
if not os.path.exists(self.data_dir):
|
||
return {}
|
||
|
||
for source in os.listdir(self.data_dir):
|
||
source_dir = os.path.join(self.data_dir, source)
|
||
if not os.path.isdir(source_dir):
|
||
continue
|
||
|
||
result[source] = {}
|
||
for i in range(days):
|
||
date_str = (datetime.now() - timedelta(days=i)).strftime("%Y-%m-%d")
|
||
file_path = os.path.join(source_dir, f"{date_str}.json")
|
||
|
||
if os.path.exists(file_path):
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
try:
|
||
data = json.load(f)
|
||
result[source][date_str] = data
|
||
except:
|
||
continue
|
||
|
||
return result
|
||
|
||
def get_saved_dates(self, source_id: str) -> List[str]:
|
||
"""
|
||
获取指定热榜源已保存的日期列表
|
||
|
||
Args:
|
||
source_id: 热榜源 ID
|
||
|
||
Returns:
|
||
已保存的日期列表（降序）
|
||
"""
|
||
dates = []
|
||
source_dir = os.path.join(self.data_dir, source_id)
|
||
|
||
if not os.path.exists(source_dir):
|
||
return []
|
||
|
||
for filename in os.listdir(source_dir):
|
||
if filename.endswith('.json'):
|
||
date_str = filename.replace('.json', '')
|
||
dates.append(date_str)
|
||
|
||
return sorted(dates, reverse=True)
|
||
|
||
def get_old_data_files(self, days: int = 7) -> List[Dict[str, str]]:
|
||
"""
|
||
获取指定天数之前的旧数据文件列表
|
||
|
||
Args:
|
||
days: 天数阈值（默认7天）
|
||
|
||
Returns:
|
||
旧文件列表 [{source_id, date_str, file_path}]
|
||
"""
|
||
from datetime import datetime, timedelta
|
||
|
||
old_files = []
|
||
cutoff_date = datetime.now() - timedelta(days=days)
|
||
|
||
if not os.path.exists(self.data_dir):
|
||
return []
|
||
|
||
for source in os.listdir(self.data_dir):
|
||
source_dir = os.path.join(self.data_dir, source)
|
||
if not os.path.isdir(source_dir):
|
||
continue
|
||
|
||
for filename in os.listdir(source_dir):
|
||
if not filename.endswith('.json'):
|
||
continue
|
||
|
||
try:
|
||
file_date = datetime.strptime(filename.replace('.json', ''), "%Y-%m-%d")
|
||
if file_date < cutoff_date:
|
||
old_files.append({
|
||
"source_id": source,
|
||
"date_str": filename.replace('.json', ''),
|
||
"file_path": os.path.join(source_dir, filename)
|
||
})
|
||
except:
|
||
continue
|
||
|
||
return old_files
|
||
|
||
def cleanup_old_files(self, files: List[Dict[str, str]]) -> int:
|
||
"""
|
||
删除指定的旧文件
|
||
|
||
Args:
|
||
files: 旧文件列表
|
||
|
||
Returns:
|
||
删除的文件数量
|
||
"""
|
||
deleted_count = 0
|
||
for f in files:
|
||
try:
|
||
if os.path.exists(f["file_path"]):
|
||
os.remove(f["file_path"])
|
||
deleted_count += 1
|
||
except Exception as e:
|
||
print(f"[DataStorage] 删除失败 {f['file_path']}: {e}")
|
||
|
||
return deleted_count
|
||
|
||
def load_hot_list(self, source_id: str, date_str: str = "") -> Optional[List[Dict]]:
|
||
"""
|
||
加载指定日期的热榜数据
|
||
|
||
Args:
|
||
source_id: 热榜源 ID
|
||
date_str: 日期（默认今天）
|
||
|
||
Returns:
|
||
热榜数据列表
|
||
"""
|
||
if not date_str:
|
||
date_str = datetime.now().strftime("%Y-%m-%d")
|
||
|
||
file_path = config.get_data_path(source_id, date_str)
|
||
|
||
if not os.path.exists(file_path):
|
||
return None
|
||
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
try:
|
||
data = json.load(f)
|
||
# 返回当天的最后一条记录
|
||
if isinstance(data, list) and len(data) > 0:
|
||
return data[-1].get("data", [])
|
||
except:
|
||
pass
|
||
|
||
return None
|
||
|
||
def list_all_data(self) -> Dict[str, int]:
|
||
"""
|
||
列出所有已保存的热榜数据统计
|
||
|
||
Returns:
|
||
源ID -> 保存记录数
|
||
"""
|
||
stats = {}
|
||
|
||
if not os.path.exists(self.data_dir):
|
||
return {}
|
||
|
||
for source in os.listdir(self.data_dir):
|
||
source_dir = os.path.join(self.data_dir, source)
|
||
if not os.path.isdir(source_dir):
|
||
continue
|
||
|
||
count = 0
|
||
for filename in os.listdir(source_dir):
|
||
if filename.endswith('.json'):
|
||
count += 1
|
||
|
||
if count > 0:
|
||
stats[source] = count
|
||
|
||
return stats
|
||
|
||
def clear_old_data(self, keep_days: int = 30) -> int:
|
||
"""
|
||
清理旧数据
|
||
|
||
Args:
|
||
keep_days: 保留最近几天的数据
|
||
|
||
Returns:
|
||
删除的文件数量
|
||
"""
|
||
deleted_count = 0
|
||
cutoff_date = datetime.now() - timedelta(days=keep_days)
|
||
|
||
if not os.path.exists(self.data_dir):
|
||
return 0
|
||
|
||
for source in os.listdir(self.data_dir):
|
||
source_dir = os.path.join(self.data_dir, source)
|
||
if not os.path.isdir(source_dir):
|
||
continue
|
||
|
||
for filename in os.listdir(source_dir):
|
||
if not filename.endswith('.json'):
|
||
continue
|
||
|
||
try:
|
||
file_date = datetime.strptime(filename.replace('.json', ''), "%Y-%m-%d")
|
||
if file_date < cutoff_date:
|
||
file_path = os.path.join(source_dir, filename)
|
||
os.remove(file_path)
|
||
deleted_count += 1
|
||
except:
|
||
continue
|
||
|
||
return deleted_count
|
||
|
||
|
||
# 全局存储实例
|
||
storage = DataStorage()
|
||
|
||
|
||
def auto_save_hot_list(source_id: str, data: Dict[str, Any]) -> bool:
|
||
"""
|
||
自动保存热榜数据的便捷函数
|
||
|
||
在获取热榜数据后调用此函数保存数据
|
||
"""
|
||
return storage.save_hot_list(source_id, data)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 测试存储功能
|
||
print("📊 每日热榜数据存储管理")
|
||
print("=" * 50)
|
||
|
||
# 列出所有数据
|
||
stats = storage.list_all_data()
|
||
if stats:
|
||
print("\n已保存的热榜数据：")
|
||
for source, count in stats.items():
|
||
print(f" • {source}: {count} 条记录")
|
||
else:
|
||
print("\n暂无保存的数据")
|
||
|
||
# 列出已保存的日期
|
||
print("\n已保存的日期（微博）：")
|
||
dates = storage.get_saved_dates("weibo")
|
||
if dates:
|
||
for date in dates[:7]:
|
||
print(f" • {date}")
|
||
else:
|
||
print(" 暂无数据")
|