NDOJ/judge/caching.py
2023-10-11 20:33:48 -05:00

80 lines
2.3 KiB
Python

from inspect import signature
from django.core.cache import cache, caches
from django.db.models.query import QuerySet
from django.core.handlers.wsgi import WSGIRequest
import hashlib
MAX_NUM_CHAR = 50
NONE_RESULT = "__None__"
def arg_to_str(arg):
if hasattr(arg, "id"):
return str(arg.id)
if isinstance(arg, list) or isinstance(arg, QuerySet):
return hashlib.sha1(str(list(arg)).encode()).hexdigest()[:MAX_NUM_CHAR]
if len(str(arg)) > MAX_NUM_CHAR:
return str(arg)[:MAX_NUM_CHAR]
return str(arg)
def filter_args(args_list):
return [x for x in args_list if not isinstance(x, WSGIRequest)]
l0_cache = caches["l0"] if "l0" in caches else None
def cache_wrapper(prefix, timeout=None):
def get_key(func, *args, **kwargs):
args_list = list(args)
signature_args = list(signature(func).parameters.keys())
args_list += [kwargs.get(k) for k in signature_args[len(args) :]]
args_list = filter_args(args_list)
args_list = [arg_to_str(i) for i in args_list]
key = prefix + ":" + ":".join(args_list)
key = key.replace(" ", "_")
return key
def _get(key):
if not l0_cache:
return cache.get(key)
print("GET", key, l0_cache.get(key))
return l0_cache.get(key) or cache.get(key)
def _set_l0(key, value):
if l0_cache:
print("SET", key, value)
l0_cache.set(key, value, 30)
def _set(key, value, timeout):
_set_l0(key, value)
cache.set(key, value, timeout)
def decorator(func):
def wrapper(*args, **kwargs):
cache_key = get_key(func, *args, **kwargs)
result = _get(cache_key)
if result is not None:
_set_l0(cache_key, result)
if result == NONE_RESULT:
result = None
return result
result = func(*args, **kwargs)
if result is None:
result = NONE_RESULT
_set(cache_key, result, timeout)
return result
def dirty(*args, **kwargs):
cache_key = get_key(func, *args, **kwargs)
cache.delete(cache_key)
if l0_cache:
l0_cache.delete(cache_key)
wrapper.dirty = dirty
return wrapper
return decorator