Repository: kmyk/online-judge-tools
Pull request #470 (2019-07-06 21:40:40)

#312: add --mle option for test subcommand
#312 でやり残していた `--mle` オプションの追加をします。
わざわざ指定するのは面倒なので、この `oj` を直接利用するユーザにはあまり嬉しさはないのですが、上部の高品質な wrapper 経由で叩くユーザにとっては嬉しいはずです。


Patch:
diff --git a/onlinejudge/_implementation/main.py b/onlinejudge/_implementation/main.py
index fa15b0d..89daea4 100644
--- a/onlinejudge/_implementation/main.py
+++ b/onlinejudge/_implementation/main.py
@@ -143,6 +143,7 @@ tips:
     subparser.add_argument('-s', '--silent', action='store_true', help='don\'t report output and correct answer even if not AC  (for --mode all)')
     subparser.add_argument('-e', '--error', type=float, help='check as floating point number: correct if its absolute or relative error doesn\'t exceed it')
     subparser.add_argument('-t', '--tle', type=float, help='set the time limit (in second) (default: inf)')
+    subparser.add_argument('--mle', type=float, help='set the memory limit (in megabyte) (default: inf)')
     subparser.add_argument('-i', '--print-input', action='store_true', help='print input cases if not AC')
     subparser.add_argument('-j', '--jobs', metavar='N', type=int, help='specifies the number of jobs to run simultaneously  (default: no parallelization)')
     subparser.add_argument('--print-memory', action='store_true', help='print the amount of memory which your program used, even if it is small enough')
diff --git a/onlinejudge/_implementation/utils.py b/onlinejudge/_implementation/utils.py
index e1de20d..92275b7 100644
--- a/onlinejudge/_implementation/utils.py
+++ b/onlinejudge/_implementation/utils.py
@@ -162,11 +162,13 @@ def exec_command(command_str: str, *, stdin: IO[Any], timeout: Optional[float] =
             answer = None
 
         end = time.perf_counter()
+        memory = None  # type: Optional[float]
         if gnu_time is not None:
             with open(fh.name) as fh1:
-                memory = int(fh1.read().splitlines()[-1]) / 1000  # type: Optional[float]
-        else:
-            memory = None
+                reported = fh1.read()
+            log.debug('GNU time says:\n%s', reported)
+            if reported.strip() and reported.splitlines()[-1].isdigit():
+                memory = int(reported.splitlines()[-1]) / 1000
     info = {
         'answer': answer,  # Optional[byte]
         'elapsed': end - begin,  # float, in second


Test patch:
diff --git a/onlinejudge/_implementation/command/test.py b/onlinejudge/_implementation/command/test.py
index 36481bb..15f2457 100644
--- a/onlinejudge/_implementation/command/test.py
+++ b/onlinejudge/_implementation/command/test.py
@@ -48,7 +48,7 @@ def compare_as_floats(xs_: str, ys_: str, error: float) -> bool:
     return True
 
 
-def compare_and_report(proc: subprocess.Popen, answer: str, test_input_path: pathlib.Path, test_output_path: Optional[pathlib.Path], *, mode: str, error: Optional[float], does_print_input: bool, silent: bool, rstrip: bool) -> str:
+def compare_and_report(proc: subprocess.Popen, answer: str, elapsed: float, memory: Optional[float], test_input_path: pathlib.Path, test_output_path: Optional[pathlib.Path], *, mle: Optional[float], mode: str, error: Optional[float], does_print_input: bool, silent: bool, rstrip: bool) -> str:
     # prepare the comparing function
     if error:  # float mode
         match = lambda a, b: compare_as_floats(a, b, error)
@@ -79,6 +79,10 @@ def compare_and_report(proc: subprocess.Popen, answer: str, test_input_path: pat
         log.failure(log.red('TLE'))
         status = 'TLE'
         print_input()
+    elif memory is not None and mle is not None and memory > mle:
+        log.failure(log.red('MLE'))
+        status = 'MLE'
+        print_input()
     elif proc.returncode != 0:
         log.failure(log.red('RE') + ': return code %d', proc.returncode)
         status = 'RE'
@@ -156,7 +160,7 @@ def test_single_case(test_name: str, test_input_path: pathlib.Path, test_output_
             else:
                 log.warning('memory: %f MB', memory)
 
-        status = compare_and_report(proc, answer, test_input_path, test_output_path, mode=args.mode, error=args.error, does_print_input=args.print_input, silent=args.silent, rstrip=args.rstrip)
+        status = compare_and_report(proc, answer, elapsed, memory, test_input_path, test_output_path, mle=args.mle, mode=args.mode, error=args.error, does_print_input=args.print_input, silent=args.silent, rstrip=args.rstrip)
 
     # return the result
     testcase = {
@@ -205,6 +209,8 @@ def test(args: 'argparse.Namespace') -> None:
     if not check_gnu_time(args.gnu_time):
         log.warning('GNU time is not available: %s', args.gnu_time)
         args.gnu_time = None
+    if args.mle is not None and args.gnu_time is None:
+        raise RuntimeError('--mle is used but GNU time does not exist')
 
     # run tests
     history = []  # type: List[Dict[str, Any]]
diff --git a/tests/command_test.py b/tests/command_test.py
index 61c7ac1..a6f40c3 100644
--- a/tests/command_test.py
+++ b/tests/command_test.py
@@ -8,8 +8,8 @@ from tests.utils import cat
 
 
 class TestTest(unittest.TestCase):
-    def snippet_call_test(self, args, files, expected):
-        result = tests.utils.run_in_sandbox(args=['-v', 'test', '--json'] + args, files=files)
+    def snippet_call_test(self, args, files, expected, verbose=True):
+        result = tests.utils.run_in_sandbox(args=(['-v'] if verbose else []) + ['test', '--json'] + args, files=files)
         self.assertTrue(result['proc'].stdout)
         data = json.loads(result['proc'].stdout.decode())
         if expected is None:
@@ -434,6 +434,7 @@ class TestTest(unittest.TestCase):
             args=['--jobs', str(PARALLEL), '--silent', '-c', tests.utils.python_c("import time; time.sleep(1); print(1)")],
             files=files,
             expected=expected,
+            verbose=False,
         )
 
     @unittest.skipIf(os.name == 'nt', "memory checking is disabled on Windows environment")
@@ -471,6 +472,28 @@ class TestTest(unittest.TestCase):
             self.assertEqual(case['status'], 'AC')
             self.assertLess(case['memory'], 100)
 
+    @unittest.skipIf(os.name == 'nt', "memory checking is disabled on Windows environment")
+    def test_call_test_memory_limit_error(self):
+        # make a bytes of 100 MB
+        self.snippet_call_test(
+            args=['--mle', '50', '-c', tests.utils.python_c("print(len(b'A' * 100000000))")],
+            files=[
+                {
+                    'path': 'test/sample-1.in',
+                    'data': 'foo\n'
+                },
+            ],
+            expected=[{
+                'status': 'MLE',
+                'testcase': {
+                    'name': 'sample-1',
+                    'input': '%s/test/sample-1.in',
+                },
+                'output': '100000000\n',
+                'exitcode': 0,
+            }],
+        )
+
     def test_call_stderr(self):
         data = self.snippet_call_test(
             args=['-c', tests.utils.python_c("import sys; print('foo', file=sys.stderr)")],