Repository: kmyk/online-judge-tools
Pull request #503 (2019-08-23 07:20:08)

NONE: add generate-input command
ランダムケースを生成するときに使うサブコマンドを追加します。
たとえば、ランダムケースを生成するスクリプトと、愚直解と、バグってる実装とを与えると、いい感じに撃墜ケースを生成してくれます。

この機能それ自体にさほど価値はありません。多少便利にはなりますが、同等のことは `or i in {000..099} ; do ./generate.py > test/random-$i.in ; done` とか `seq 100 | parallel './generate.py | tee test/random-{}.in | ./a.out > test/random-{}.out'` などでできて、ほぼ困らないためです。
これはむしろドキュメントとして働かせることを意図しています。「ランダムケースを生成してテストをする」という発想をまだ持っていない競プロerはおそらく多いはずですが、具体的にサブコマンドとして提供することでこの発想を自然に導入できます。

(宛先は `develop`)


Patch:
diff --git a/onlinejudge/_implementation/command/generate_input.py b/onlinejudge/_implementation/command/generate_input.py
new file mode 100644
index 0000000..b733379
--- /dev/null
+++ b/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)
diff --git a/onlinejudge/_implementation/main.py b/onlinejudge/_implementation/main.py
index 86ca1c5..853f6b6 100644
--- a/onlinejudge/_implementation/main.py
+++ b/onlinejudge/_implementation/main.py
@@ -13,6 +13,7 @@ import onlinejudge.__about__ as version
 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 @@ tips:
     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':
diff --git a/onlinejudge/_implementation/utils.py b/onlinejudge/_implementation/utils.py
index 722857c..2704f8a 100644
--- a/onlinejudge/_implementation/utils.py
+++ b/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
diff --git a/readme.md b/readme.md
index 13890a1..2b11e72 100644
--- a/readme.md
+++ b/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
 


Test patch:
diff --git a/tests/command_generate_input.py b/tests/command_generate_input.py
new file mode 100644
index 0000000..278f04f
--- /dev/null
+++ b/tests/command_generate_input.py
@@ -0,0 +1,115 @@
+import os
+import unittest
+
+import tests.utils
+
+
+class GenerateInputTest(unittest.TestCase):
+    def snippet_call_generate_input(self, args, input_files, expected_values, disallowed_files=None):
+        with tests.utils.sandbox(input_files) as tempdir:
+            tests.utils.run(['generate-input'] + args, check=True)
+            for expect in expected_values:
+                self.assertTrue(os.path.exists(expect['path']))
+                if expect['data'] is not None:
+                    with open(expect['path']) as f:
+                        self.assertEqual(''.join(f.readlines()), expect['data'])
+            if disallowed_files is not None:
+                for file in disallowed_files:
+                    self.assertFalse(os.path.exists(file))
+
+    def test_call_generate_input_parallel(self):
+        self.snippet_call_generate_input(
+            args=[tests.utils.python_script('generate.py'), '-j', '4'],
+            input_files=[
+                {
+                    'path': 'generate.py',
+                    'data': 'print("Hello, world!")\n'
+                },
+            ],
+            expected_values=[{
+                'path': 'test/random-{}.in'.format(str(i).zfill(3)),
+                'data': 'Hello, world!\n'
+            } for i in range(100)],
+            disallowed_files=[
+                'test/random-100.in',
+            ],
+        )
+
+    def test_call_generate_input_with_output(self):
+        expected_values = []
+        for i in range(20):
+            expected_values += [{
+                'path': 'test/random-{}.in'.format(str(i).zfill(3)),
+                'data': None,
+            }]
+            expected_values += [{
+                'path': 'test/random-{}.out'.format(str(i).zfill(3)),
+                'data': None,
+            }]
+        self.snippet_call_generate_input(
+            args=[tests.utils.python_script('generate.py'), '20', '--command', tests.utils.python_script('main.py')],
+            input_files=[
+                {
+                    'path': 'generate.py',
+                    'data': 'import random\nprint(random.randrange(100))\n'
+                },
+                {
+                    'path': 'main.py',
+                    'data': 'print(int(input()) * 2)\n'
+                },
+            ],
+            expected_values=expected_values,
+            disallowed_files=[
+                'test/random-100.in',
+            ],
+        )
+
+    def test_call_generate_input_hack(self):
+        self.snippet_call_generate_input(args=[tests.utils.python_script('generate.py'), '--command', tests.utils.python_script('ac.py'), '--hack', tests.utils.python_script('wa.py')], input_files=[
+            {
+                'path': 'generate.py',
+                'data': 'import random\nprint(random.randint(1, 10))\n'
+            },
+            {
+                'path': 'ac.py',
+                'data': 'print(int(input()) // 10)\n'
+            },
+            {
+                'path': 'wa.py',
+                'data': 'print(0)\n'
+            },
+        ], expected_values=[
+            {
+                'path': 'test/hack-000.in',
+                'data': '10\n'
+            },
+            {
+                'path': 'test/hack-000.out',
+                'data': '1\n'
+            },
+        ])
+
+    def test_call_generate_input_hack_parallel(self):
+        self.snippet_call_generate_input(args=[tests.utils.python_script('generate.py'), '--command', tests.utils.python_script('ac.py'), '--hack', tests.utils.python_script('wa.py'), '-j', '4'], input_files=[
+            {
+                'path': 'generate.py',
+                'data': 'import random\nprint(random.randint(1, 100))\n'
+            },
+            {
+                'path': 'ac.py',
+                'data': 'print(int(input()) // 100)\n'
+            },
+            {
+                'path': 'wa.py',
+                'data': 'print(0)\n'
+            },
+        ], expected_values=[
+            {
+                'path': 'test/hack-000.in',
+                'data': '100\n'
+            },
+            {
+                'path': 'test/hack-000.out',
+                'data': '1\n'
+            },
+        ])
diff --git a/tests/utils.py b/tests/utils.py
index ff2c8c0..ad6959c 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -68,3 +68,8 @@ def cat():
 def python_c(cmd):
     assert '"' not in cmd
     return '{} -c "{}"'.format(sys.executable, cmd)
+
+
+def python_script(path):
+    assert '"' not in path
+    return '{} "{}"'.format(sys.executable, path)