Cache Problem Testcases

judge/utils/problems.py

@@ -1,6 +1,8 @@
 from collections import defaultdict
 from math import e
+import os, zipfile
 
+from django.conf import settings
 from django.core.cache import cache
 from django.db.models import Case, Count, ExpressionWrapper, F, Max, Q, When
 from django.db.models.fields import FloatField
@@ -145,3 +147,61 @@ def hot_problems(duration, limit):
 
         cache.set(cache_key, qs, 900)
     return qs
+
+
+def get_visible_content(data):
+    data = data or b''
+    data = data.replace(b'\r\n', b'\r').replace(b'\r', b'\n')
+
+    data = data.decode('utf-8')
+
+    if (len(data) > settings.TESTCASE_VISIBLE_LENGTH):
+        data = data[:settings.TESTCASE_VISIBLE_LENGTH]
+        data += '.' * 3
+    return data
+
+
+def get_cachekey_file(file):
+    return file.replace(' ', '===')
+
+def get_problem_case(problem, files):
+    result = {}
+    uncached_files = []
+
+    for file in files:
+        cache_key = 'problem_archive:%s:%s' % (problem.code, get_cachekey_file(file))
+        qs = cache.get(cache_key)
+        if qs is None:
+            uncached_files.append(file)
+        else:
+            result['file'] = qs
+
+    if not uncached_files:
+        return result
+
+    archive_path = os.path.join(settings.DMOJ_PROBLEM_DATA_ROOT,
+                                    str(problem.data_files.zipfile))
+    if not os.path.exists(archive_path):
+        raise Exception(
+            'archive file "%s" does not exist' % archive_path)
+    try:
+        archive = zipfile.ZipFile(archive_path, 'r')
+    except zipfile.BadZipfile:
+        raise Exception('bad archive: "%s"' % archive_path)
+
+    for file in uncached_files:
+        cache_key = 'problem_archive:%s:%s' % (problem.code, get_cachekey_file(file))
+        with archive.open(file) as f:
+            s = f.read(settings.TESTCASE_VISIBLE_LENGTH + 3)
+            # add this so there are no characters left behind (ex, 'á' = 2 utf-8 chars)
+            while True:
+                try:
+                    s.decode('utf-8')
+                    break
+                except UnicodeDecodeError:
+                    s += f.read(1)
+            qs = get_visible_content(s)
+        cache.set(cache_key, qs, 86400)
+        result[file] = qs
+
+    return result