#318: fix about the feature for "Test Cases" of AtCoderSubmission

onlinejudge/implementation/utils.py

@@ -259,7 +259,10 @@ def remove_suffix(s: str, suffix: str) -> str:
 tzinfo_jst = datetime.timezone(datetime.timedelta(hours=+9), 'JST')
 
 
-def getter_with_load_details(name: str) -> Callable:
+def getter_with_load_details(name: str, check_with: Optional[str] = None) -> Callable:
+    if check_with is None:
+        check_with = name
+
     @functools.wraps(lambda self: getattr(self, name))
     def wrapper(self, session: Optional[requests.Session] = None):
         if getattr(self, name) is None:

onlinejudge/service/atcoder.py

@@ -562,8 +562,10 @@ def _load_details(self, session: Optional[requests.Session] = None) -> None:
         self._score = int(data['Score'])
         self._code_size = int(utils.remove_suffix(data['Code Size'], ' Byte'))
         self._status = data['Status']
-        self._exec_time_msec = int(utils.remove_suffix(data['Exec Time'], ' ms'))
-        self._memory_byte = int(utils.remove_suffix(data['Memory'], ' KB')) * 1000
+        if 'Exec Time' in data:
+            self._exec_time_msec = int(utils.remove_suffix(data['Exec Time'], ' ms'))
+        if 'Memory' in data:
+            self._memory_byte = int(utils.remove_suffix(data['Memory'], ' KB')) * 1000
 
         # Compile Error
         compile_error = soup.find('h4', text='Compile Error')
@@ -591,26 +593,31 @@ def get_problem(self, session: Optional[requests.Session] = None) -> AtCoderProb
     get_score = utils.getter_with_load_details('_score')  # type: Callable[..., int]
     get_code_size = utils.getter_with_load_details('_code_size')  # type: Callable[..., int]
     get_status = utils.getter_with_load_details('_status')  # type: Callable[..., str]
-    get_exec_time_msec = utils.getter_with_load_details('_exec_time_msec')  # type: Callable[..., int]
-    get_memory_byte = utils.getter_with_load_details('_memory_byte')  # type: Callable[..., int]
+    get_exec_time_msec = utils.getter_with_load_details('_exec_time_msec', check_with='_status')  # type: Callable[..., int]
+    get_memory_byte = utils.getter_with_load_details('_memory_byte', check_with='_status')  # type: Callable[..., int]
     get_test_cases = utils.getter_with_load_details('_test_cases')  # type: Callable[..., List[AtCoderSubmissionTestCaseResult]]
 
 
 class AtCoderSubmissionTestCaseResult(object):
-    def __init__(self, case_name: str, status: str, exec_time_msec: int, memory_kb: int):
+    def __init__(self, case_name: str, status: str, exec_time_msec: Optional[int], memory_byte: Optional[int]):
         self.case_name = case_name
         self.status = status
         self.exec_time_msec = exec_time_msec
-        self.memory_kb = memory_kb
+        self.memory_byte = memory_byte
 
     @classmethod
     def _from_table_row(cls, tr: bs4.Tag) -> 'AtCoderSubmissionTestCaseResult':
         tds = tr.find_all('td')
         case_name = tds[0].text
         status = tds[1].text
-        exec_time_msec = int(utils.remove_suffix(tds[2].text, ' ms'))
-        memory_kb = int(utils.remove_suffix(tds[3].text, ' KB'))
-        return AtCoderSubmissionTestCaseResult(case_name, status, exec_time_msec, memory_kb)
+        exec_time_msec = None  # type: Optional[int]
+        memory_byte = None  # type: Optional[int]
+        if len(tds) == 4:
+            exec_time_msec = int(utils.remove_suffix(tds[2].text, ' ms'))
+            memory_byte = int(utils.remove_suffix(tds[3].text, ' KB')) * 1000
+        else:
+            assert len(tds) == 2
+        return AtCoderSubmissionTestCaseResult(case_name, status, exec_time_msec, memory_byte)
 
 
 onlinejudge.dispatch.services += [AtCoderService]

tests/service_atcoder.py

@@ -76,6 +76,19 @@ def test_submission_info(self):
         self.assertEqual(submission.get_exec_time_msec(), 85)
         self.assertEqual(submission.get_memory_byte(), 3328 * 1000)
 
+    def test_get_test_cases(self):
+        submission = AtCoderSubmission.from_url('https://atcoder.jp/contests/tricky/submissions/119944')
+        test_cases = submission.get_test_cases()
+        self.assertEqual(len(test_cases), 2)
+        self.assertEqual(test_cases[0].case_name, 'input_01.txt')
+        self.assertEqual(test_cases[0].status, 'TLE')
+        self.assertEqual(test_cases[0].exec_time_msec, None)
+        self.assertEqual(test_cases[0].memory_byte, None)
+        self.assertEqual(test_cases[1].case_name, 'input_02.txt')
+        self.assertEqual(test_cases[1].status, 'AC')
+        self.assertEqual(test_cases[1].exec_time_msec, 131)
+        self.assertEqual(test_cases[1].memory_byte, 7400 * 1000)
+
     def test_get_source_code(self):
         submission = AtCoderSubmission.from_url('https://atcoder.jp/contests/abc100/submissions/3082514')
         self.assertEqual(submission.get_source_code(), b'/9\\|\\B/c:(\ncYay!')