#312: add a feature to check MLE

onlinejudge/_implementation/command/generate_output.py

@@ -25,9 +25,9 @@ def generate_output_single_case(test_name: str, test_input_path: pathlib.Path, *
 
     # run the command
     with test_input_path.open() as inf:
-        begin = time.perf_counter()
-        answer, proc = utils.exec_command(args.command, stdin=inf, timeout=args.tle)
-        end = time.perf_counter()
+        info, proc = utils.exec_command(args.command, stdin=inf, timeout=args.tle)
+        answer = info['answer']  # type: Optional[bytes]
+        elapsed = info['elapsed']  # type: float
 
     # acquire lock to print logs properly, if in parallel
     nullcontext = contextlib.ExitStack()
@@ -37,7 +37,7 @@ def generate_output_single_case(test_name: str, test_input_path: pathlib.Path, *
             log.info('%s', test_name)
 
         # check the result
-        log.status('time: %f sec', end - begin)
+        log.status('time: %f sec', elapsed)
         if proc.returncode is None:
             log.failure(log.red('TLE'))
             log.info('skipped.')
@@ -46,6 +46,7 @@ def generate_output_single_case(test_name: str, test_input_path: pathlib.Path, *
             log.failure(log.red('RE') + ': return code %d', proc.returncode)
             log.info('skipped.')
             return
+        assert answer is not None
         log.emit(utils.snip_large_file_content(answer, limit=40, head=20, tail=10, bold=True))
 
         # find the destination path

onlinejudge/_implementation/command/test.py

@@ -6,8 +6,10 @@
 import pathlib
 import subprocess
 import sys
+import tempfile
 import threading
 import time
+import traceback
 from typing import *
 
 import onlinejudge
@@ -18,6 +20,9 @@
 if TYPE_CHECKING:
     import argparse
 
+MEMORY_WARNING = 500  # megabyte
+MEMORY_PRINT = 100  # megabyte
+
 
 def compare_as_floats(xs_: str, ys_: str, error: float) -> bool:
     def f(x):
@@ -129,12 +134,11 @@ def test_single_case(test_name: str, test_input_path: pathlib.Path, test_output_
 
     # run the binary
     with test_input_path.open() as inf:
-        begin = time.perf_counter()
-        answer_byte, proc = utils.exec_command(args.command, stdin=inf, timeout=args.tle)
-        end = time.perf_counter()
-        elapsed = end - begin
-        answer = answer_byte.decode()  # TODO: the `answer` should be bytes, not str
-        proc.terminate()
+        info, proc = utils.exec_command(args.command, stdin=inf, timeout=args.tle, gnu_time=args.gnu_time)
+        # TODO: the `answer` should be bytes, not str
+        answer = (info['answer'] or b'').decode()  # type: str
+        elapsed = info['elapsed']  # type: float
+        memory = info['memory']  # type: Optional[float]
 
     # lock is require to avoid mixing logs if in parallel
     nullcontext = contextlib.ExitStack()  # TODO: use contextlib.nullcontext() after updating Python to 3.7
@@ -143,6 +147,14 @@ def test_single_case(test_name: str, test_input_path: pathlib.Path, test_output_
             log.emit('')
             log.info('%s', test_name)
         log.status('time: %f sec', elapsed)
+        if memory:
+            if memory < MEMORY_PRINT:
+                if args.print_memory:
+                    log.status('memory: %f MB', memory)
+            elif memory < MEMORY_WARNING:
+                log.status('memory: %f MB', memory)
+            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)
 
@@ -153,14 +165,32 @@ def test_single_case(test_name: str, test_input_path: pathlib.Path, test_output_
     }
     if test_output_path:
         testcase['output'] = str(test_output_path.resolve())
-    result = {
+    return {
         'status': status,
         'testcase': testcase,
         'output': answer,
         'exitcode': proc.returncode,
         'elapsed': elapsed,
+        'memory': memory,
     }
-    return result
+
+
+def check_gnu_time(gnu_time: str) -> bool:
+    try:
+        with tempfile.NamedTemporaryFile(delete=True) as fh:
+            proc = subprocess.run([gnu_time, '-f', '%M KB', '-o', fh.name, '--quiet', '--', 'true'])
+            assert proc.returncode == 0
+            with open(fh.name) as fh1:
+                data = fh1.read()
+            int(utils.remove_suffix(data.rstrip(), ' KB'))
+            return True
+    except NameError:
+        raise  # NameError is not a runtime error caused by the environmet, but a coding mistake
+    except AttributeError:
+        raise  # AttributeError is also a mistake
+    except Exception as e:
+        log.debug(traceback.format_exc())
+    return False
 
 
 def test(args: 'argparse.Namespace') -> None:
@@ -171,6 +201,11 @@ def test(args: 'argparse.Namespace') -> None:
         args.test = fmtutils.drop_backup_or_hidden_files(args.test)
     tests = fmtutils.construct_relationship_of_files(args.test, args.directory, args.format)
 
+    # check wheather GNU time is available
+    if not check_gnu_time(args.gnu_time):
+        log.warning('GNU time is not available: %s', args.gnu_time)
+        args.gnu_time = None
+
     # run tests
     history = []  # type: List[Dict[str, Any]]
     if args.jobs is None:
@@ -186,19 +221,29 @@ def test(args: 'argparse.Namespace') -> None:
                 history += [future.result()]
 
     # summarize
-    slowest = -1  # type: Union[int, float]
+    slowest = -1.0  # type: float
     slowest_name = ''
+    heaviest = -1.0  # type: float
+    heaviest_name = ''
     ac_count = 0
     for result in history:
         if result['status'] == 'AC':
             ac_count += 1
         if slowest < result['elapsed']:
             slowest = result['elapsed']
             slowest_name = result['testcase']['name']
+        if result['memory'] is not None and heaviest < result['memory']:
+            heaviest = result['memory']
+            heaviest_name = result['testcase']['name']
 
     # print the summary
     log.emit('')
     log.status('slowest: %f sec  (for %s)', slowest, slowest_name)
+    if heaviest >= 0:
+        if heaviest < MEMORY_WARNING:
+            log.status('max memory: %f MB  (for %s)', heaviest, heaviest_name)
+        else:
+            log.warning('max memory: %f MB  (for %s)', heaviest, heaviest_name)
     if ac_count == len(tests):
         log.success('test ' + log.green('success') + ': %d cases', len(tests))
     else:

onlinejudge/_implementation/main.py

@@ -142,6 +142,8 @@ def get_parser() -> argparse.ArgumentParser:
     subparser.add_argument('-t', '--tle', type=float, help='set the time limit (in second) (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')
+    subparser.add_argument('--gnu-time', help='used to measure memory consumption (default: "time")', default='time')
     subparser.add_argument('--no-ignore-backup', action='store_false', dest='ignore_backup')
     subparser.add_argument('--ignore-backup', action='store_true', help='ignore backup files and hidden files (i.e. files like "*~", "\\#*\\#" and ".*") (default)')
     subparser.add_argument('--json', action='store_true')

onlinejudge/_implementation/utils.py

@@ -13,6 +13,7 @@
 import shutil
 import subprocess
 import sys
+import tempfile
 import time
 import urllib.parse
 from typing import *
@@ -134,21 +135,43 @@ def textfile(s: str) -> str:  # should have trailing newline
         return s + '\n'
 
 
-def exec_command(command: str, timeout: Optional[float] = None, **kwargs) -> Tuple[bytes, subprocess.Popen]:
-    try:
-        proc = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE, stderr=sys.stderr, **kwargs)
-    except FileNotFoundError:
-        log.error('No such file or directory: %s', command)
-        sys.exit(1)
-    except PermissionError:
-        log.error('Permission denied: %s', command)
-        sys.exit(1)
-    try:
-        answer, _ = proc.communicate(timeout=timeout)
-    except subprocess.TimeoutExpired:
-        proc.terminate()
-        answer = b''
-    return answer, proc
+def exec_command(command_str: str, *, stdin: IO[Any], timeout: Optional[float] = None, gnu_time: Optional[str] = None) -> Tuple[Dict[str, Any], subprocess.Popen]:
+    if gnu_time is not None:
+        context = tempfile.NamedTemporaryFile(delete=True)  # type: Any
+    else:
+        context = contextlib.ExitStack()  # TODO: we should use contextlib.nullcontext() if possible
+    with context as fh:
+        command = shlex.split(command_str)
+        if gnu_time is not None:
+            command = [gnu_time, '-f', '%M', '-o', fh.name, '--quiet', '--'] + command
+        begin = time.perf_counter()
+
+        try:
+            proc = subprocess.Popen(command, stdin=stdin, stdout=subprocess.PIPE, stderr=sys.stderr)
+        except FileNotFoundError:
+            log.error('No such file or directory: %s', command)
+            sys.exit(1)
+        except PermissionError:
+            log.error('Permission denied: %s', command)
+            sys.exit(1)
+        try:
+            answer, _ = proc.communicate(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            proc.terminate()
+            answer = None
+
+        end = time.perf_counter()
+        if gnu_time is not None:
+            with open(fh.name) as fh1:
+                memory = int(fh1.read()) / 1000  # type: Optional[float]
+        else:
+            memory = None
+    info = {
+        'answer': answer,  # Optional[byte]
+        'elapsed': end - begin,  # float, in second
+        'memory': memory,  # Optional[float], in megabyte
+    }
+    return info, proc
 
 
 # We should use this instead of posixpath.normpath

tests/command_test.py

@@ -455,3 +455,36 @@ def test_call_test_in_parallel(self):
             files=files,
             expected=expected,
         )
+
+    def test_call_test_large_memory(self):
+        # make a bytes of 100 MB
+        data = self.snippet_call_test(
+            args=['-c', """python -c 'print(len(b"A" * 100000000))'"""],
+            files=[
+                {
+                    'path': 'test/sample-1.in',
+                    'data': 'foo\n'
+                },
+            ],
+            expected=None,
+        )
+        for case in data:
+            self.assertEqual(case['status'], 'AC')
+            self.assertGreater(case['memory'], 100)
+            self.assertLess(case['memory'], 1000)
+
+    def test_call_test_small_memory(self):
+        # just print "foo"
+        data = self.snippet_call_test(
+            args=['-c', """python -c 'print("foo")'"""],
+            files=[
+                {
+                    'path': 'test/sample-1.in',
+                    'data': 'foo\n'
+                },
+            ],
+            expected=None,
+        )
+        for case in data:
+            self.assertEqual(case['status'], 'AC')
+            self.assertLess(case['memory'], 100)