Add memory unit to problem admin

judge/admin/problem.py

@@ -33,11 +33,14 @@ from judge.widgets import (
     HeavyPreviewPageDownWidget,
 )
 
+MEMORY_UNITS = (("KB", "KB"), ("MB", "MB"))
+
 
 class ProblemForm(ModelForm):
     change_message = forms.CharField(
         max_length=256, label="Edit reason", required=False
     )
+    memory_unit = forms.ChoiceField(choices=MEMORY_UNITS)
 
     def __init__(self, *args, **kwargs):
         super(ProblemForm, self).__init__(*args, **kwargs)
@@ -51,6 +54,12 @@ class ProblemForm(ModelForm):
             }
         )
 
+    def clean(self):
+        memory_unit = self.cleaned_data.get("memory_unit", "KB")
+        if memory_unit == "MB":
+            self.cleaned_data["memory_limit"] *= 1024
+        return self.cleaned_data
+
     class Meta:
         widgets = {
             "authors": AdminHeavySelect2MultipleWidget(
@@ -94,13 +103,28 @@ class ProblemCreatorListFilter(admin.SimpleListFilter):
 
 
 class LanguageLimitInlineForm(ModelForm):
+    memory_unit = forms.ChoiceField(choices=MEMORY_UNITS, label=_("Memory unit"))
+
     class Meta:
-        widgets = {"language": AdminSelect2Widget}
+        widgets = {
+            "language": AdminSelect2Widget,
+            "memory_limit": TextInput(attrs={"size": "10"}),
+        }
+
+    def clean(self):
+        if not self.cleaned_data.get("language"):
+            self.cleaned_data["DELETE"] = True
+        if (
+            self.cleaned_data.get("memory_limit")
+            and self.cleaned_data.get("memory_unit") == "MB"
+        ):
+            self.cleaned_data["memory_limit"] *= 1024
+        return self.cleaned_data
 
 
 class LanguageLimitInline(admin.TabularInline):
     model = LanguageLimit
-    fields = ("language", "time_limit", "memory_limit")
+    fields = ("language", "time_limit", "memory_limit", "memory_unit")
     form = LanguageLimitInlineForm
 
 
@@ -204,7 +228,7 @@ class ProblemAdmin(CompareVersionAdmin):
         ),
         (_("Taxonomy"), {"fields": ("types", "group")}),
         (_("Points"), {"fields": (("points", "partial"), "short_circuit")}),
-        (_("Limits"), {"fields": ("time_limit", "memory_limit")}),
+        (_("Limits"), {"fields": ("time_limit", ("memory_limit", "memory_unit"))}),
         (_("Language"), {"fields": ("allowed_languages",)}),
         (_("Justice"), {"fields": ("banned_users",)}),
         (_("History"), {"fields": ("change_message",)}),

judge/migrations/0130_auto_20220831_1048.py

@@ -8,18 +8,24 @@ import judge.utils.problem_data
 class Migration(migrations.Migration):
 
     dependencies = [
-        ('judge', '0129_auto_20220622_1424'),
+        ("judge", "0129_auto_20220622_1424"),
     ]
 
     operations = [
         migrations.AddField(
-            model_name='problem',
-            name='pdf_description',
-            field=models.FileField(blank=True, null=True, storage=judge.utils.problem_data.ProblemDataStorage(), upload_to=judge.models.problem.problem_directory_file, verbose_name='pdf statement'),
+            model_name="problem",
+            name="pdf_description",
+            field=models.FileField(
+                blank=True,
+                null=True,
+                storage=judge.utils.problem_data.ProblemDataStorage(),
+                upload_to=judge.models.problem.problem_directory_file,
+                verbose_name="pdf statement",
+            ),
         ),
         migrations.AlterField(
-            model_name='problem',
-            name='description',
-            field=models.TextField(blank=True, verbose_name='problem body'),
+            model_name="problem",
+            name="description",
+            field=models.TextField(blank=True, verbose_name="problem body"),
         ),
     ]

judge/models/problem.py

@@ -18,7 +18,10 @@ from judge.models.profile import Organization, Profile
 from judge.models.runtime import Language
 from judge.user_translations import gettext as user_gettext
 from judge.utils.raw_sql import RawSQLColumn, unique_together_left_join
-from judge.models.problem_data import problem_data_storage, problem_directory_file_helper
+from judge.models.problem_data import (
+    problem_data_storage,
+    problem_directory_file_helper,
+)
 
 __all__ = [
     "ProblemGroup",
@@ -32,9 +35,11 @@ __all__ = [
     "TranslatedProblemForeignKeyQuerySet",
 ]
 
+
 def problem_directory_file(data, filename):
     return problem_directory_file_helper(data.code, filename)
 
+
 class ProblemType(models.Model):
     name = models.CharField(
         max_length=20, verbose_name=_("problem category ID"), unique=True

judge/views/problem.py

@@ -417,9 +417,7 @@ class ProblemPdfDescriptionView(ProblemMixin, SingleObjectMixin, View):
                 response.content = f.read()
 
         response["Content-Type"] = "application/pdf"
-        response["Content-Disposition"] = "inline; filename=%s.pdf" % (
-            problem.code,
-        )
+        response["Content-Disposition"] = "inline; filename=%s.pdf" % (problem.code,)
         return response