File size: 4,894 Bytes
a226682 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import os
import json
import hashlib
import pickle
from typing import Any, Optional, Dict # 添加 Dict
from pathlib import Path
from config import Config
class CacheManager:
"""缓存管理器"""
def __init__(self, cache_dir: str = None):
self.cache_dir = cache_dir or Config.CACHE_DIR
Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
def _get_cache_key(self, *args, **kwargs) -> str:
"""生成缓存键"""
content = str(args) + str(sorted(kwargs.items()))
return hashlib.md5(content.encode()).hexdigest()
def _get_cache_path(self, key: str) -> str:
"""获取缓存文件路径"""
return os.path.join(self.cache_dir, f"{key}.pkl")
def get(self, key: str) -> Optional[Any]:
"""获取缓存"""
if not Config.ENABLE_CACHE:
return None
cache_file = self._get_cache_path(key)
if os.path.exists(cache_file):
try:
with open(cache_file, 'rb') as f:
return pickle.load(f)
except Exception as e:
print(f"读取缓存失败 ({key}): {e}")
return None
return None
def set(self, key: str, value: Any) -> bool:
"""设置缓存"""
if not Config.ENABLE_CACHE:
return False
cache_file = self._get_cache_path(key)
try:
with open(cache_file, 'wb') as f:
pickle.dump(value, f)
return True
except Exception as e:
print(f"缓存保存失败 ({key}): {e}")
return False
def exists(self, key: str) -> bool:
"""检查缓存是否存在"""
if not Config.ENABLE_CACHE:
return False
return os.path.exists(self._get_cache_path(key))
def delete(self, key: str) -> bool:
"""删除指定缓存"""
cache_file = self._get_cache_path(key)
if os.path.exists(cache_file):
try:
os.remove(cache_file)
return True
except Exception as e:
print(f"删除缓存失败 ({key}): {e}")
return False
return False
def clear(self, pattern: str = None):
"""清除缓存"""
if not os.path.exists(self.cache_dir):
return
count = 0
for file in os.listdir(self.cache_dir):
if file.endswith('.pkl'):
if pattern is None or pattern in file:
try:
os.remove(os.path.join(self.cache_dir, file))
count += 1
except Exception as e:
print(f"删除缓存文件失败 ({file}): {e}")
print(f"已清除 {count} 个缓存文件")
def get_cache_size(self) -> int:
"""获取缓存总大小（字节）"""
if not os.path.exists(self.cache_dir):
return 0
total_size = 0
for file in os.listdir(self.cache_dir):
if file.endswith('.pkl'):
file_path = os.path.join(self.cache_dir, file)
total_size += os.path.getsize(file_path)
return total_size
def get_cache_info(self) -> dict:
"""获取缓存信息"""
if not os.path.exists(self.cache_dir):
return {
'count': 0,
'size': 0,
'size_mb': 0
}
count = 0
total_size = 0
for file in os.listdir(self.cache_dir):
if file.endswith('.pkl'):
count += 1
file_path = os.path.join(self.cache_dir, file)
total_size += os.path.getsize(file_path)
return {
'count': count,
'size': total_size,
'size_mb': round(total_size / (1024 * 1024), 2)
}
def save_json(self, key: str, data: dict) -> bool:
"""保存为JSON格式（用于可读性）"""
cache_file = os.path.join(self.cache_dir, f"{key}.json")
try:
with open(cache_file, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=2)
return True
except Exception as e:
print(f"保存JSON缓存失败 ({key}): {e}")
return False
def load_json(self, key: str) -> Optional[dict]:
"""加载JSON格式缓存"""
cache_file = os.path.join(self.cache_dir, f"{key}.json")
if os.path.exists(cache_file):
try:
with open(cache_file, 'r', encoding='utf-8') as f:
return json.load(f)
except Exception as e:
print(f"读取JSON缓存失败 ({key}): {e}")
return None
return None |