[teamai] Push 1 resource(s) from root
This commit is contained in:
@@ -0,0 +1,345 @@
|
||||
# -*- 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(" 暂无数据")
|
||||
Reference in New Issue
Block a user