Skip to content

Instantly share code, notes, and snippets.

@Ryaang

Ryaang/cache_on_disk.py

Last active January 13, 2025 03:29
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save Ryaang/88693c067941f4628a3d3a7bb10f3db1 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save Ryaang/88693c067941f4628a3d3a7bb10f3db1 to your computer and use it in GitHub Desktop.
Download ZIP
支持同步和异步函数的持久化缓存装饰器，用于缓存计算结果，避免重复计算。缓存结果存储在指定目录中，后续调用时直接读取。自动适配同步和异步函数。
Raw
cache_on_disk.py
from diskcache import Cache
import functools
import asyncio
from typing import Any, Callable, TypeVar, Union, Awaitable
import time
import logging
import hashlib
import json
T = TypeVar('T')
def persistent_cache(
cache_dir: str = '.cache',
expire: int = None, # 过期时间(秒)
tag: str = None, # 缓存标签
max_size: int = 2 * 1024 * 1024 * 1024, # 2GB
eviction_policy: str = 'least-recently-used',
verbose: bool = False
):
"""持久化缓存装饰器,支持同步和异步函数
Args:
cache_dir: 缓存文件存储目录
expire: 缓存过期时间(秒)
tag: 缓存标签,用于批量管理缓存
"""
# 创建缓存实例
cache = Cache(cache_dir, size_limit=max_size)
def decorator(func: Callable[..., Union[T, Awaitable[T]]]) -> Callable[..., Union[T, Awaitable[T]]]:
"""实际的装饰器函数"""
def make_cache_key(*args, **kwargs):
"""生成缓存键的哈希值"""
def serialize_arg(arg):
"""序列化参数"""
if isinstance(arg, (dict, list, set)):
return json.dumps(arg, sort_keys=True)
return str(arg)
# 构建键的组成部分
key_parts = [
func.__name__, # 函数名
tuple(serialize_arg(arg) for arg in args), # 位置参数
frozenset((k, serialize_arg(v)) for k, v in sorted(kwargs.items())) # 关键字参数
]
# 将键转换为字符串并编码
key_str = json.dumps(str(key_parts), sort_keys=True)
# 使用 SHA-256 生成哈希
hash_obj = hashlib.sha256(key_str.encode('utf-8'))
return hash_obj.hexdigest()
@functools.wraps(func)
def sync_wrapper(*args, **kwargs) -> T:
"""同步函数的包装器"""
# 使用函数名和参数作为键
key = make_cache_key(*args, **kwargs)
# 尝试获取缓存
result = cache.get(key)
if result is not None:
if verbose:
logging.debug(f"Cache hit for: {func.__name__}")
return result
# 执行函数并缓存结果
result = func(*args, **kwargs)
cache.set(key, result, expire=expire, tag=tag)
return result
@functools.wraps(func)
async def async_wrapper(*args, **kwargs) -> T:
"""异步函数的包装器"""
key = make_cache_key(*args, **kwargs)
try:
result = cache.get(key)
if result is not None:
if verbose:
logging.debug(f"Cache hit for: {func.__name__} (key: {key[:8]}...)")
return result
result = await func(*args, **kwargs)
try:
cache.set(key, result, expire=expire, tag=tag)
except Exception as e:
logging.warning(f"Failed to cache result for key {key[:8]}...: {e}")
return result
except Exception as e:
logging.error(f"Cache error for key {key[:8]}...: {e}")
return await func(*args, **kwargs)
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
return decorator
if __name__ == "__main__":
# 异步函数
@persistent_cache(cache_dir='.pcache')
async def expensive_function(a: str, b: str) -> str:
return a + b
# 运行并计时
import time
start_time = time.time()
result = asyncio.run(expensive_function("a"*10000, "b"*10000))
end_time = time.time()
print(f"Time taken: {end_time - start_time} seconds")
# print(result)
@Ryaang
Copy link
Author

Ryaang commented Jan 13, 2025

持久化缓存装饰器

这是一个支持同步和异步函数的持久化缓存装饰器，用于缓存函数计算结果，避免重复计算，提升性能。缓存结果会被存储在指定目录中，后续调用时直接从缓存读取。

核心功能：

  1. 自动缓存：根据函数名和参数哈希值生成唯一的缓存文件。
  2. 同步/异步支持：自动适配同步和异步函数。
  3. 持久化存储：缓存结果存储在磁盘中，程序重启后依然有效。
  4. 参数哈希：通过参数的哈希值确保不同参数组合对应不同的缓存。

使用方法：

  1. 使用 @persistent_cache(cache_dir='.pcache') 装饰需要缓存的函数。
  2. 对于异步函数，直接使用 await 调用；对于同步函数，正常调用即可。

示例：

@persistent_cache(cache_dir='.pcache')
async def expensive_function(a: str, b: str) -> str:
    return a + b

# 第一次调用会计算结果并缓存
result = await expensive_function("a"*10000, "b"*10000)

# 第二次调用会直接从缓存读取
result = await expensive_function("a"*10000, "b"*10000)

适用场景：

  • 计算成本高的函数（如复杂计算、网络请求、数据库查询）。
  • 输入参数不变时输出结果不变的函数。

注意事项：

  • 缓存文件会占用磁盘空间，需定期清理。
  • 不适用于依赖外部状态（如时间、全局变量）的函数。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment