Add cache wrapper

judge/caching.py

@@ -1,10 +1,48 @@
+from inspect import signature
 from django.core.cache import cache
+from django.db.models.query import QuerySet
+
+import hashlib
+
+MAX_NUM_CHAR = 15
 
 
-def finished_submission(sub):
-    keys = ["user_complete:%d" % sub.user_id, "user_attempted:%s" % sub.user_id]
-    if hasattr(sub, "contest"):
-        participation = sub.contest.participation
-        keys += ["contest_complete:%d" % participation.id]
-        keys += ["contest_attempted:%d" % participation.id]
-    cache.delete_many(keys)
+def cache_wrapper(prefix, timeout=86400):
+    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 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 = [arg_to_str(i) for i in args_list]
+        key = prefix + ":" + ":".join(args_list)
+        key = key.replace(" ", "_")
+        return key
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            cache_key = get_key(func, *args, **kwargs)
+            result = cache.get(cache_key)
+            if result is not None:
+                return result
+
+            result = func(*args, **kwargs)
+            cache.set(cache_key, result, timeout)
+            return result
+
+        def dirty(*args, **kwargs):
+            cache_key = get_key(func, *args, **kwargs)
+            cache.delete(cache_key)
+
+        wrapper.dirty = dirty
+
+        return wrapper
+
+    return decorator