From 729e3d5a89fcb3d90d349509b6892397ab03a66f Mon Sep 17 00:00:00 2001
From: Tudor Brindus <me@tbrindus.ca>
Date: Fri, 6 Mar 2020 22:48:35 -0500
Subject: [PATCH] Fix a couple of code style issues

---
 dmoj/judge.py   | 27 ++++++++-------------------
 dmoj/problem.py | 13 +++++++++++++
 dmoj/result.py  |  4 ++--
 3 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/dmoj/judge.py b/dmoj/judge.py
index e645ae7a6..9851fa991 100644
--- a/dmoj/judge.py
+++ b/dmoj/judge.py
@@ -9,7 +9,7 @@
 from http.server import HTTPServer
 from itertools import chain
 
-from dmoj import graders, packet
+from dmoj import packet
 from dmoj.control import JudgeControlRequestHandler
 from dmoj.error import CompileError
 from dmoj.judgeenv import clear_problem_dirs_cache, env, get_supported_problems, startup_warnings
@@ -82,17 +82,8 @@ def update_problems(self):
         self.updater_signal.set()
 
     def _block_and_grade(self, problem, language, source, short_circuit, report=print):
-        if 'signature_grader' in problem.config:
-            grader_class = graders.SignatureGrader
-        elif 'interactive' in problem.config:
-            grader_class = graders.BridgedInteractiveGrader
-        elif 'custom_judge' in problem.config:
-            grader_class = graders.CustomGrader
-        else:
-            grader_class = graders.StandardGrader
-
         try:
-            self.current_grader = grader_class(self, problem, language, utf8bytes(source))
+            self.current_grader = problem.grader_class(self, problem, language, utf8bytes(source))
         except CompileError as compilation_error:
             error = compilation_error.args[0] or b'compiler exited abnormally'
 
@@ -187,25 +178,23 @@ def grade_cases(self, grader, cases, short_circuit=False, is_short_circuiting=Fa
             if isinstance(case, BatchedTestCase):
                 yield BatchBegin()
 
-                for batched_case in self.grade_cases(grader, case.batched_cases,
-                                                     short_circuit=case.config['short_circuit'],
-                                                     is_short_circuiting=is_short_circuiting):
+                for batched_result in self.grade_cases(grader, case.batched_cases,
+                                                       short_circuit=case.config['short_circuit'],
+                                                       is_short_circuiting=is_short_circuiting):
                     # A batched case just failed.
                     # There are two cases where this means that we should completely short-circuit:
                     # 1. If the batch was worth 0 points, to emulate the property of 0-point cases.
                     # 2. If the short_circuit flag is true, see <https://github.com/DMOJ/judge/issues/341>.
-                    if (batched_case.result_flag & Result.WA) and \
+                    if (batched_result.result_flag & Result.WA) and \
                             (short_circuit or case.kind in (TestCase.KIND_SAMPLE, TestCase.KIND_PRETEST)):
                         is_short_circuiting = True
-                    yield batched_case
+                    yield batched_result
                 yield BatchEnd()
                 continue
 
             # Stop grading if we're short circuiting
             if is_short_circuiting:
-                result = Result(case)
-                result.result_flag = Result.SC
-                yield result
+                yield Result(case, Result.SC)
                 continue
 
             result = grader.grade(case)
diff --git a/dmoj/problem.py b/dmoj/problem.py
index a122e5879..794426450 100644
--- a/dmoj/problem.py
+++ b/dmoj/problem.py
@@ -165,6 +165,19 @@ def _resolve_archive_files(self):
             return archive
         return None
 
+    @property
+    def grader_class(self):
+        from dmoj import graders
+
+        if 'signature_grader' in self.config:
+            return graders.SignatureGrader
+        elif 'interactive' in self.config:
+            return graders.BridgedInteractiveGrader
+        elif 'custom_judge' in self.config:
+            return graders.CustomGrader
+        else:
+            return graders.StandardGrader
+
 
 class ProblemDataManager(dict):
     def __init__(self, problem_id, **kwargs):
diff --git a/dmoj/result.py b/dmoj/result.py
index 98a9b2e52..37cc29f54 100644
--- a/dmoj/result.py
+++ b/dmoj/result.py
@@ -25,8 +25,8 @@ class Result:
         'IE': 'red',
     }
 
-    def __init__(self, case):
-        self.result_flag = 0
+    def __init__(self, case, result_flag=0):
+        self.result_flag = result_flag
         self.execution_time = 0
         self.wall_clock_time = 0
         self.max_memory = 0