NONE: add generate-input command

onlinejudge/_implementation/command/generate_input.py

@@ -0,0 +1,231 @@
+# Python Version: 3.x
+import argparse
+import concurrent.futures
+import contextlib
+import itertools
+import os
+import pathlib
+import threading
+from typing import *
+
+import onlinejudge._implementation.format_utils as fmtutils
+import onlinejudge._implementation.logging as log
+import onlinejudge._implementation.utils as utils
+
+
+@contextlib.contextmanager
+def BufferedExecutor(lock: Optional[threading.Lock]):
+    buf = []  # type: List[Tuple[Callable, List[Any], Dict[str, Any]]]
+
+    def submit(f, *args, **kwargs):
+        nonlocal buf
+        if lock is None:
+            f(*args, **kwargs)
+        else:
+            buf += [(f, args, kwargs)]
+
+    result = yield submit
+
+    if lock is not None:
+        with lock:
+            for f, args, kwargs in buf:
+                f(*args, **kwargs)
+    return result
+
+
+def write_result(input_data: bytes, output_data: Optional[bytes], *, input_path: pathlib.Path, output_path: pathlib.Path, print_data: bool, lock: Optional[threading.Lock] = None) -> None:
+    # acquire lock to print logs properly, if in parallel
+    nullcontext = contextlib.ExitStack()
+    with lock or nullcontext:
+
+        if not input_path.parent.is_dir():
+            os.makedirs(str(input_path.parent), exist_ok=True)
+
+        if print_data:
+            log.emit('input:')
+            log.emit(utils.snip_large_file_content(input_data, limit=40, head=20, tail=10, bold=True))
+        with input_path.open('wb') as fh:
+            fh.write(input_data)
+        log.success('saved to: %s', input_path)
+
+        if output_data is not None:
+            if print_data:
+                log.emit('output:')
+                log.emit(utils.snip_large_file_content(output_data, limit=40, head=20, tail=10, bold=True))
+            with output_path.open('wb') as fh:
+                fh.write(output_data)
+            log.success('saved to: %s', output_path)
+
+
+def check_status(info, proc, *, submit):
+    submit(log.status, 'time: %f sec', info['elapsed'])
+    if proc.returncode is None:
+        submit(log.failure, log.red('TLE'))
+        submit(log.info, 'skipped.')
+        return False
+    elif proc.returncode != 0:
+        submit(log.failure, log.red('RE') + ': return code %d', proc.returncode)
+        submit(log.info, 'skipped.')
+        return False
+    assert info['answer'] is not None
+    return True
+
+
+def generate_input_single_case(generator: str, *, input_path: pathlib.Path, output_path: pathlib.Path, command: Optional[str], tle: Optional[float], name: str, lock: Optional[threading.Lock] = None) -> None:
+    with BufferedExecutor(lock) as submit:
+
+        # print the header
+        submit(log.emit, '')
+        submit(log.info, '%s', name)
+
+        # generate input
+        submit(log.status, 'generate input...')
+        info, proc = utils.exec_command(generator, timeout=tle)
+        input_data = info['answer']  # type: bytes
+        if not check_status(info, proc, submit=submit):
+            return None
+
+        # generate output
+        if command is None:
+            output_data = None  # type: Optional[bytes]
+        else:
+            submit(log.status, 'generate output...')
+            info, proc = utils.exec_command(command, input=input_data, timeout=tle)
+            output_data = info['answer']
+            if not check_status(info, proc, submit=submit):
+                return None
+
+        # write result
+        submit(write_result, input_data=input_data, output_data=output_data, input_path=input_path, output_path=output_path, print_data=True)
+
+
+def simple_match(a: str, b: str) -> bool:
+    if a == b:
+        return True
+    if a.rstrip() == b.rstrip():
+        log.warning('WA if no rstrip')
+        return True
+    return False
+
+
+def try_hack_once(generator: str, command: str, hack: str, *, tle: Optional[float], attempt: int, lock: Optional[threading.Lock] = None) -> Optional[Tuple[bytes, bytes]]:
+    with BufferedExecutor(lock) as submit:
+
+        # print the header
+        submit(log.emit, '')
+        submit(log.info, '%d-th attempt', attempt)
+
+        # generate input
+        submit(log.status, 'generate input...')
+        info, proc = utils.exec_command(generator, stdin=None, timeout=tle)
+        input_data = info['answer']  # type: Optional[bytes]
+        if not check_status(info, proc, submit=submit):
+            return None
+        assert input_data is not None
+
+        # generate output
+        submit(log.status, 'generate output...')
+        info, proc = utils.exec_command(command, input=input_data, timeout=tle)
+        output_data = info['answer']  # type: Optional[bytes]
+        if not check_status(info, proc, submit=submit):
+            return None
+        assert output_data is not None
+
+        # hack
+        submit(log.status, 'hack...')
+        info, proc = utils.exec_command(hack, input=input_data, timeout=tle)
+        answer = (info['answer'] or b'').decode()  # type: str
+        elapsed = info['elapsed']  # type: float
+        memory = info['memory']  # type: Optional[float]
+
+        # compare
+        status = 'AC'
+        if proc.returncode is None:
+            submit(log.failure, log.red('TLE'))
+            status = 'TLE'
+        elif proc.returncode != 0:
+            log.failure(log.red('RE') + ': return code %d', proc.returncode)
+            status = 'RE'
+        expected = output_data.decode()
+        if not simple_match(answer, expected):
+            log.failure(log.red('WA'))
+            log.emit('input:\n%s', utils.snip_large_file_content(input_data, limit=40, head=20, tail=10, bold=True))
+            log.emit('output:\n%s', utils.snip_large_file_content(answer.encode(), limit=40, head=20, tail=10, bold=True))
+            log.emit('expected:\n%s', utils.snip_large_file_content(output_data, limit=40, head=20, tail=10, bold=True))
+            status = 'WA'
+
+        if status == 'AC':
+            return None
+        else:
+            return (input_data, output_data)
+
+
+def generate_input(args: argparse.Namespace) -> None:
+    if args.hack and not args.command:
+        raise RuntimeError('--hack must be used with --command')
+
+    if args.name is None:
+        if args.hack:
+            args.name = 'hack'
+        else:
+            args.name = 'random'
+
+    if args.count is None:
+        if args.hack:
+            args.count = 1
+        else:
+            args.count = 100
+
+    def iterate_path():
+        for i in itertools.count():
+            name = '{}-{}'.format(args.name, str(i).zfill(args.width))
+            input_path = fmtutils.path_from_format(args.directory, args.format, name=name, ext='in')
+            output_path = fmtutils.path_from_format(args.directory, args.format, name=name, ext='out')
+            if not input_path.exists() and not output_path.exists():
+                yield (name, input_path, output_path)
+
+    # generate cases
+    if args.jobs is None:
+        for name, input_path, output_path in itertools.islice(iterate_path(), args.count):
+            if not args.hack:
+                # generate serially
+                generate_input_single_case(args.generator, input_path=input_path, output_path=output_path, command=args.command, tle=args.tle, name=name)
+
+            else:
+                # hack serially
+                for attempt in itertools.count(1):
+                    data = try_hack_once(args.generator, command=args.command, hack=args.hack, tle=args.tle, attempt=attempt)
+                    if data is not None:
+                        write_result(*data, input_path=input_path, output_path=output_path, print_data=False)
+                        break
+    else:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=args.jobs) as executor:
+            lock = threading.Lock()
+            futures = []  # type: List[concurrent.futures.Future]
+
+            if not args.hack:
+                # generate concurrently
+                for name, input_path, output_path in itertools.islice(iterate_path(), args.count):
+                    futures += [executor.submit(generate_input_single_case, args.generator, input_path=input_path, output_path=output_path, command=args.command, tle=args.tle, name=name, lock=lock)]
+                for future in futures:
+                    future.result()
+
+            else:
+                # hack concurrently
+                attempt = 0
+                for _ in range(args.jobs):
+                    attempt += 1
+                    futures += [executor.submit(try_hack_once, args.generator, command=args.command, hack=args.hack, tle=args.tle, attempt=attempt, lock=lock)]
+                for _, input_path, output_path in itertools.islice(iterate_path(), args.count):
+                    data = None
+                    while data is None:
+                        concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_COMPLETED)
+                        for i in range(len(futures)):
+                            if not futures[i].done():
+                                continue
+                            data = futures[i].result()
+                            attempt += 1
+                            futures[i] = executor.submit(try_hack_once, args.generator, command=args.command, hack=args.hack, tle=args.tle, attempt=attempt, lock=lock)
+                            if data is not None:
+                                break
+                    write_result(*data, input_path=input_path, output_path=output_path, print_data=False, lock=lock)

onlinejudge/_implementation/main.py

@@ -13,6 +13,7 @@
 import onlinejudge._implementation.logging as log
 import onlinejudge._implementation.utils as utils
 from onlinejudge._implementation.command.download import download
+from onlinejudge._implementation.command.generate_input import generate_input
 from onlinejudge._implementation.command.generate_output import generate_output
 from onlinejudge._implementation.command.get_standings import get_standings
 from onlinejudge._implementation.command.login import login
@@ -172,6 +173,27 @@ def get_parser() -> argparse.ArgumentParser:
     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)')
 
+    # generate input
+    subparser = subparsers.add_parser('generate-input', aliases=['g/i'], help='generate input files form given generator', formatter_class=argparse.RawTextHelpFormatter, epilog='''\
+format string for --format:
+  %s                    name
+  %e                    extension: "in" or "out"
+  (both %d and %e are required.)
+
+tips:
+  You can do similar things with shell: e.g. `for i in {000..099} ; do ./generate.py > test/random-$i.in ; done`
+''')
+    subparser.add_argument('-f', '--format', default='%s.%e', help='a format string to recognize the relationship of test cases. (default: "%%s.%%e")')
+    subparser.add_argument('-d', '--directory', type=pathlib.Path, default=pathlib.Path('test'), help='a directory name for test cases (default: test/)')
+    subparser.add_argument('-t', '--tle', type=float, help='set the time limit (in second) (default: inf)')
+    subparser.add_argument('-j', '--jobs', type=int, help='run tests in parallel')
+    subparser.add_argument('--width', type=int, default=3, help='specify the width of indices of cases. (default: 3)')
+    subparser.add_argument('--name', help='specify the base name of cases. (default: "random")')
+    subparser.add_argument('-c', '--command', help='specify your solution to generate output')
+    subparser.add_argument('--hack', help='specify your solution to be compared the reference solution given by --command')
+    subparser.add_argument('generator', type=str, help='your program to generate test cases')
+    subparser.add_argument('count', nargs='?', type=int, help='the number of cases to generate (default: 100)')
+
     # split input
     subparser = subparsers.add_parser('split-input', help='split a input file which contains many cases, using your implementation  (experimental)', formatter_class=argparse.RawTextHelpFormatter, epilog='''\
 format string for --output:
@@ -269,6 +291,8 @@ def run_program(args: argparse.Namespace, parser: argparse.ArgumentParser) -> No
         test_reactive(args)
     elif args.subcommand in ['generate-output', 'g/o']:
         generate_output(args)
+    elif args.subcommand in ['generate-input', 'g/i']:
+        generate_input(args)
     elif args.subcommand == 'split-input':
         split_input(args)
     elif args.subcommand == 'get-standings':

onlinejudge/_implementation/utils.py

@@ -134,7 +134,10 @@ def textfile(s: str) -> str:  # should have trailing newline
         return s + '\n'
 
 
-def exec_command(command_str: str, *, stdin: IO[Any], timeout: Optional[float] = None, gnu_time: Optional[str] = None) -> Tuple[Dict[str, Any], subprocess.Popen]:
+def exec_command(command_str: str, *, stdin: Optional[IO[Any]] = None, input: Optional[bytes] = None, timeout: Optional[float] = None, gnu_time: Optional[str] = None) -> Tuple[Dict[str, Any], subprocess.Popen]:
+    if input is not None:
+        assert stdin is None
+        stdin = subprocess.PIPE  # type: ignore
     if gnu_time is not None:
         context = tempfile.NamedTemporaryFile(delete=True)  # type: Any
     else:
@@ -157,7 +160,7 @@ def exec_command(command_str: str, *, stdin: IO[Any], timeout: Optional[float] =
             log.error('Permission denied: %s', command)
             sys.exit(1)
         try:
-            answer, _ = proc.communicate(timeout=timeout)
+            answer, _ = proc.communicate(input=input, timeout=timeout)
         except subprocess.TimeoutExpired:
             proc.terminate()
             answer = None

readme.md

@@ -52,6 +52,7 @@ Tools for online judge services. Downloading sample cases, Testing/Submitting yo
     -   yukicoder
 -   Test your solution
 -   Test your solution for reactive problem
+-   Generate input files from generators
 -   Generate output files from input and reference implementation
 -   Split an input file with many cases to files