From 80f838b9a86f419b43edbc155beae421f4dffd42 Mon Sep 17 00:00:00 2001
From: Mateusz Masiarz <matimasiarz@gmail.com>
Date: Sun, 11 Feb 2024 16:22:03 +0100
Subject: [PATCH] Create ocen archive on `export` (#179)

* Add ocen archive creation

* Fancy ingen printing

* Package for tests

* Add tests

* Fix tests

* Bump version
---
 src/sinol_make/__init__.py                  |  2 +-
 src/sinol_make/commands/export/__init__.py  | 98 +++++++++++++++++----
 src/sinol_make/commands/ingen/ingen_util.py |  2 +
 src/sinol_make/commands/outgen/__init__.py  | 10 ++-
 tests/commands/export/test_integration.py   | 36 ++++++++
 tests/commands/export/test_unit.py          |  8 ++
 tests/packages/ocen/config.yml              |  4 +
 tests/packages/ocen/in/ocen0.in             |  1 +
 tests/packages/ocen/in/ocen0a.in            |  1 +
 tests/packages/ocen/in/ocen1a.in            |  1 +
 tests/packages/ocen/in/ocen1ocen.in         |  1 +
 tests/packages/ocen/out/ocen0.out           |  1 +
 tests/packages/ocen/prog/ocen.cpp           |  9 ++
 tests/packages/ocen/prog/oceningen.cpp      | 15 ++++
 tests/util.py                               |  8 ++
 15 files changed, 177 insertions(+), 20 deletions(-)
 create mode 100644 tests/packages/ocen/config.yml
 create mode 100644 tests/packages/ocen/in/ocen0.in
 create mode 100644 tests/packages/ocen/in/ocen0a.in
 create mode 100644 tests/packages/ocen/in/ocen1a.in
 create mode 100644 tests/packages/ocen/in/ocen1ocen.in
 create mode 100644 tests/packages/ocen/out/ocen0.out
 create mode 100644 tests/packages/ocen/prog/ocen.cpp
 create mode 100644 tests/packages/ocen/prog/oceningen.cpp

diff --git a/src/sinol_make/__init__.py b/src/sinol_make/__init__.py
index c7511fe8..4ec5b19f 100644
--- a/src/sinol_make/__init__.py
+++ b/src/sinol_make/__init__.py
@@ -9,7 +9,7 @@
 from sinol_make import util, oiejq
 
 
-__version__ = "1.5.22"
+__version__ = "1.5.23"
 
 
 def configure_parsers():
diff --git a/src/sinol_make/commands/export/__init__.py b/src/sinol_make/commands/export/__init__.py
index fd882ec5..38148f45 100644
--- a/src/sinol_make/commands/export/__init__.py
+++ b/src/sinol_make/commands/export/__init__.py
@@ -3,6 +3,7 @@
 import stat
 import shutil
 import tarfile
+import tempfile
 import argparse
 import yaml
 
@@ -10,6 +11,7 @@
 from sinol_make.commands.ingen.ingen_util import get_ingen, compile_ingen, run_ingen, ingen_exists
 from sinol_make.helpers import package_util, parsers, paths
 from sinol_make.interfaces.BaseCommand import BaseCommand
+from sinol_make.commands.outgen import Command as OutgenCommand, compile_correct_solution, get_correct_solution
 
 
 class Command(BaseCommand):
@@ -25,34 +27,91 @@ def configure_subparser(self, subparser: argparse.ArgumentParser):
             self.get_name(),
             help='Create archive for oioioi upload',
             description='Creates archive in the current directory ready to upload to sio2 or szkopul.')
+        parser.add_argument('-c', '--cpus', type=int,
+                            help=f'number of cpus to use to generate output files '
+                                 f'(default: {util.default_cpu_count()})',
+                            default=util.default_cpu_count())
         parsers.add_compilation_arguments(parser)
 
-    def get_generated_tests(self):
-        """
-        Returns list of generated tests.
-        Executes ingen to check what tests are generated.
-        """
-        if not ingen_exists(self.task_id):
-            return []
-
+    def generate_input_tests(self):
+        print('Generating tests...')
         temp_package = paths.get_cache_path('export', 'tests')
         if os.path.exists(temp_package):
             shutil.rmtree(temp_package)
         os.makedirs(temp_package)
         in_dir = os.path.join(temp_package, 'in')
-        prog_dir = os.path.join(temp_package, 'prog')
         os.makedirs(in_dir)
-        shutil.copytree(os.path.join(os.getcwd(), 'prog'), prog_dir)
+        out_dir = os.path.join(temp_package, 'out')
+        os.makedirs(out_dir)
+        prog_dir = os.path.join(temp_package, 'prog')
+        if os.path.exists(os.path.join(os.getcwd(), 'prog')):
+            shutil.copytree(os.path.join(os.getcwd(), 'prog'), prog_dir)
+
+        if ingen_exists(self.task_id):
+            ingen_path = get_ingen(self.task_id)
+            ingen_path = os.path.join(prog_dir, os.path.basename(ingen_path))
+            ingen_exe = compile_ingen(ingen_path, self.args, self.args.weak_compilation_flags)
+            if not run_ingen(ingen_exe, in_dir):
+                util.exit_with_error('Failed to run ingen.')
+
+    def generate_output_files(self):
+        tests = paths.get_cache_path('export', 'tests')
+        in_dir = os.path.join(tests, 'in')
+        os.makedirs(in_dir, exist_ok=True)
+        out_dir = os.path.join(tests, 'out')
+        os.makedirs(out_dir, exist_ok=True)
+
+        # Only example output tests are required for export.
+        ocen_tests = glob.glob(os.path.join(in_dir, f'{self.task_id}0*.in')) + \
+                     glob.glob(os.path.join(in_dir, f'{self.task_id}*ocen.in'))
+        outputs = []
+        for test in ocen_tests:
+            outputs.append(os.path.join(out_dir, os.path.basename(test).replace('.in', '.out')))
+        if len(outputs) > 0:
+            outgen = OutgenCommand()
+            correct_solution_exe = compile_correct_solution(get_correct_solution(self.task_id), self.args,
+                                                            self.args.weak_compilation_flags)
+            outgen.args = self.args
+            outgen.correct_solution_exe = correct_solution_exe
+            outgen.generate_outputs(outputs)
 
-        ingen_path = get_ingen(self.task_id)
-        ingen_path = os.path.join(prog_dir, os.path.basename(ingen_path))
-        ingen_exe = compile_ingen(ingen_path, self.args, self.args.weak_compilation_flags)
-        if not run_ingen(ingen_exe, in_dir):
-            util.exit_with_error('Failed to run ingen.')
+    def get_generated_tests(self):
+        """
+        Returns list of generated tests.
+        Executes ingen to check what tests are generated.
+        """
+        if not ingen_exists(self.task_id):
+            return []
 
+        in_dir = paths.get_cache_path('export', 'tests', 'in')
         tests = glob.glob(os.path.join(in_dir, f'{self.task_id}*.in'))
         return [package_util.extract_test_id(test, self.task_id) for test in tests]
 
+    def create_ocen(self, target_dir: str):
+        """
+        Creates ocen archive for sio2.
+        :param target_dir: Path to exported package.
+        """
+        attachments_dir = os.path.join(target_dir, 'attachments')
+        if not os.path.exists(attachments_dir):
+            os.makedirs(attachments_dir)
+        tests_dir = paths.get_cache_path('export', 'tests')
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ocen_dir = os.path.join(tmpdir, self.task_id)
+            os.makedirs(ocen_dir)
+            in_dir = os.path.join(ocen_dir, 'in')
+            os.makedirs(in_dir)
+            out_dir = os.path.join(ocen_dir, 'out')
+            os.makedirs(out_dir)
+            for ext in ['in', 'out']:
+                for test in glob.glob(os.path.join(tests_dir, ext, f'{self.task_id}0*.{ext}')) + \
+                            glob.glob(os.path.join(tests_dir, ext, f'{self.task_id}*ocen.{ext}')):
+                    shutil.copy(test, os.path.join(ocen_dir, ext, os.path.basename(test)))
+
+            with tarfile.open(os.path.join(attachments_dir, f'{self.task_id}ocen.tgz'), "w:gz") as tar:
+                tar.add(ocen_dir, arcname=os.path.basename(ocen_dir))
+
     def copy_package_required_files(self, target_dir: str):
         """
         Copies package files and directories from
@@ -79,7 +138,6 @@ def copy_package_required_files(self, target_dir: str):
             for test in glob.glob(os.path.join(os.getcwd(), ext, f'{self.task_id}0*.{ext}')):
                 shutil.copy(test, os.path.join(target_dir, ext))
 
-        print('Generating tests...')
         generated_tests = self.get_generated_tests()
         tests_to_copy = []
         for ext in ['in', 'out']:
@@ -87,11 +145,17 @@ def copy_package_required_files(self, target_dir: str):
                 if package_util.extract_test_id(test, self.task_id) not in generated_tests:
                     tests_to_copy.append((ext, test))
 
+        cache_test_dir = paths.get_cache_path('export', 'tests')
         if len(tests_to_copy) > 0:
             print(util.warning(f'Found {len(tests_to_copy)} tests that are not generated by ingen.'))
             for test in tests_to_copy:
                 print(util.warning(f'Copying {os.path.basename(test[1])}...'))
                 shutil.copy(test[1], os.path.join(target_dir, test[0], os.path.basename(test[1])))
+                shutil.copy(test[1], os.path.join(cache_test_dir, test[0], os.path.basename(test[1])))
+
+        self.generate_output_files()
+        print('Generating ocen archive...')
+        self.create_ocen(target_dir)
 
     def clear_files(self, target_dir: str):
         """
@@ -112,6 +176,7 @@ def create_makefile_in(self, target_dir: str, config: dict):
         with open(os.path.join(target_dir, 'makefile.in'), 'w') as f:
             cxx_flags = '-std=c++20'
             c_flags = '-std=gnu99'
+
             def format_multiple_arguments(obj):
                 if isinstance(obj, str):
                     return obj
@@ -163,6 +228,7 @@ def run(self, args: argparse.Namespace):
         os.makedirs(export_package_path)
 
         util.change_stack_size_to_unlimited()
+        self.generate_input_tests()
         self.copy_package_required_files(export_package_path)
         self.clear_files(export_package_path)
         self.create_makefile_in(export_package_path, config)
diff --git a/src/sinol_make/commands/ingen/ingen_util.py b/src/sinol_make/commands/ingen/ingen_util.py
index 2e3f9779..31269ed4 100644
--- a/src/sinol_make/commands/ingen/ingen_util.py
+++ b/src/sinol_make/commands/ingen/ingen_util.py
@@ -82,6 +82,7 @@ def run_ingen(ingen_exe, working_dir=None):
         st = os.stat(ingen_exe)
         os.chmod(ingen_exe, st.st_mode | stat.S_IEXEC)
 
+    print(util.bold(' Ingen output '.center(util.get_terminal_size()[1], '=')))
     process = subprocess.Popen([ingen_exe], stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                cwd=working_dir, shell=is_shell)
     while process.poll() is None:
@@ -89,5 +90,6 @@ def run_ingen(ingen_exe, working_dir=None):
 
     print(process.stdout.read().decode('utf-8'), end='')
     exit_code = process.returncode
+    print(util.bold(' End of ingen output '.center(util.get_terminal_size()[1], '=')))
 
     return exit_code == 0
diff --git a/src/sinol_make/commands/outgen/__init__.py b/src/sinol_make/commands/outgen/__init__.py
index 3ea3efb6..dfa6676c 100644
--- a/src/sinol_make/commands/outgen/__init__.py
+++ b/src/sinol_make/commands/outgen/__init__.py
@@ -38,7 +38,8 @@ def generate_outputs(self, outputs_to_generate):
         arguments = []
         for output in outputs_to_generate:
             output_basename = os.path.basename(output)
-            input = os.path.join(os.getcwd(), 'in', os.path.splitext(output_basename)[0] + '.in')
+            in_dir = os.path.join("/", *(os.path.abspath(output).split(os.sep)[:-2]), 'in')
+            input = os.path.join(in_dir, os.path.splitext(output_basename)[0] + '.in')
             arguments.append(OutputGenerationArguments(self.correct_solution_exe, input, output))
 
         with mp.Pool(self.args.cpus) as pool:
@@ -55,11 +56,14 @@ def generate_outputs(self, outputs_to_generate):
             else:
                 print(util.info('Successfully generated all output files.'))
 
-    def calculate_md5_sums(self):
+    def calculate_md5_sums(self, tests=None):
         """
         Calculates md5 sums for each test.
         :return: Tuple (dictionary of md5 sums, list of outputs tests that need to be generated)
         """
+        if tests is None:
+            tests = glob.glob(os.path.join(os.getcwd(), 'in', '*.in'))
+
         old_md5_sums = None
         try:
             with open(os.path.join(os.getcwd(), 'in', '.md5sums'), 'r') as f:
@@ -71,7 +75,7 @@ def calculate_md5_sums(self):
 
         md5_sums = {}
         outputs_to_generate = []
-        for file in glob.glob(os.path.join(os.getcwd(), 'in', '*.in')):
+        for file in tests:
             basename = os.path.basename(file)
             output_basename = os.path.splitext(os.path.basename(basename))[0] + '.out'
             output_path = os.path.join(os.getcwd(), 'out', output_basename)
diff --git a/tests/commands/export/test_integration.py b/tests/commands/export/test_integration.py
index dc020cc6..3b5d67b7 100644
--- a/tests/commands/export/test_integration.py
+++ b/tests/commands/export/test_integration.py
@@ -127,3 +127,39 @@ def test_handwritten_tests(create_package):
         extracted = os.path.join(tmpdir, task_id)
         for file in ["in/hwr0.in", "in/hwr0a.in", "out/hwr0.out", "out/hwr0a.out"]:
             assert os.path.exists(os.path.join(extracted, file))
+
+
+@pytest.mark.parametrize("create_package", [util.get_ocen_package_path()], indirect=True)
+def test_ocen_archive(create_package):
+    """
+    Test creation of ocen archive.
+    """
+    parser = configure_parsers()
+    args = parser.parse_args(["export"])
+    command = Command()
+    command.run(args)
+    task_id = package_util.get_task_id()
+    in_handwritten = ["ocen0.in", "ocen0a.in", "ocen1a.in", "ocen1ocen.in"]
+    out_handwritten = ["ocen0.out"]
+    ocen_tests = ["ocen0", "ocen0a", "ocen0b", "ocen1ocen", "ocen2ocen"]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        package_path = os.path.join(tmpdir, task_id)
+        os.mkdir(package_path)
+        with tarfile.open(f'{task_id}.tgz', "r") as tar:
+            sinol_util.extract_tar(tar, tmpdir)
+
+        for ext in ["in", "out"]:
+            tests = [os.path.basename(f) for f in glob.glob(os.path.join(package_path, ext, f'*.{ext}'))]
+            assert set(tests) == set(in_handwritten if ext == "in" else out_handwritten)
+
+        ocen_tar = os.path.join(package_path, "attachments", f"{task_id}ocen.tgz")
+        assert os.path.exists(ocen_tar)
+        ocen_dir = os.path.join(package_path, "ocen_dir")
+
+        with tarfile.open(ocen_tar, "r") as tar:
+            sinol_util.extract_tar(tar, ocen_dir)
+
+        for ext in ["in", "out"]:
+            tests = [os.path.basename(f) for f in glob.glob(os.path.join(ocen_dir, task_id, ext, f'*.{ext}'))]
+            assert set(tests) == set([f'{test}.{ext}' for test in ocen_tests])
diff --git a/tests/commands/export/test_unit.py b/tests/commands/export/test_unit.py
index 86cad2bd..4a167bd2 100644
--- a/tests/commands/export/test_unit.py
+++ b/tests/commands/export/test_unit.py
@@ -12,6 +12,10 @@ def _create_package(tmpdir, path):
     os.chdir(package_path)
     command = get_command()
     util.create_ins_outs(package_path)
+    command.args = argparse.Namespace(cpus=1, weak_compilation_flags=False,
+                                      cpp_compiler_path=compiler.get_cpp_compiler_path(),
+                                      c_compiler_path=None, python_interpreter_path=None,
+                                      java_compiler_path=None)
     return command
 
 
@@ -21,9 +25,11 @@ def test_get_generated_tests():
     """
     with tempfile.TemporaryDirectory() as tmpdir:
         command = _create_package(tmpdir, util.get_handwritten_package_path())
+        command.generate_input_tests()
         assert set(command.get_generated_tests()) == {"1a", "2a"}
 
         command = _create_package(tmpdir, util.get_simple_package_path())
+        command.generate_input_tests()
         assert set(command.get_generated_tests()) == {"1a", "2a", "3a", "4a"}
 
 
@@ -36,6 +42,7 @@ def test_copy_package_required_files():
         res_dir = os.path.join(tmpdir, "res")
         os.mkdir(res_dir)
         command = _create_package(tmpdir, util.get_handwritten_package_path())
+        command.generate_input_tests()
         command.copy_package_required_files(res_dir)
 
         assert_configs_equal(os.getcwd(), res_dir)
@@ -47,6 +54,7 @@ def test_copy_package_required_files():
         shutil.rmtree(res_dir)
         os.mkdir(res_dir)
         command = _create_package(tmpdir, util.get_simple_package_path())
+        command.generate_input_tests()
         command.copy_package_required_files(res_dir)
 
         assert_configs_equal(os.getcwd(), res_dir)
diff --git a/tests/packages/ocen/config.yml b/tests/packages/ocen/config.yml
new file mode 100644
index 00000000..c714bd75
--- /dev/null
+++ b/tests/packages/ocen/config.yml
@@ -0,0 +1,4 @@
+title: Package for testing ocen archive creation
+sinol_task_id: ocen
+time_limit: 1000
+memory_limit: 10240
diff --git a/tests/packages/ocen/in/ocen0.in b/tests/packages/ocen/in/ocen0.in
new file mode 100644
index 00000000..654d5269
--- /dev/null
+++ b/tests/packages/ocen/in/ocen0.in
@@ -0,0 +1 @@
+2 3
diff --git a/tests/packages/ocen/in/ocen0a.in b/tests/packages/ocen/in/ocen0a.in
new file mode 100644
index 00000000..ebcee1a5
--- /dev/null
+++ b/tests/packages/ocen/in/ocen0a.in
@@ -0,0 +1 @@
+3 4
diff --git a/tests/packages/ocen/in/ocen1a.in b/tests/packages/ocen/in/ocen1a.in
new file mode 100644
index 00000000..8d04f961
--- /dev/null
+++ b/tests/packages/ocen/in/ocen1a.in
@@ -0,0 +1 @@
+1 2
diff --git a/tests/packages/ocen/in/ocen1ocen.in b/tests/packages/ocen/in/ocen1ocen.in
new file mode 100644
index 00000000..99818b5e
--- /dev/null
+++ b/tests/packages/ocen/in/ocen1ocen.in
@@ -0,0 +1 @@
+3 5
diff --git a/tests/packages/ocen/out/ocen0.out b/tests/packages/ocen/out/ocen0.out
new file mode 100644
index 00000000..7ed6ff82
--- /dev/null
+++ b/tests/packages/ocen/out/ocen0.out
@@ -0,0 +1 @@
+5
diff --git a/tests/packages/ocen/prog/ocen.cpp b/tests/packages/ocen/prog/ocen.cpp
new file mode 100644
index 00000000..9d30bab0
--- /dev/null
+++ b/tests/packages/ocen/prog/ocen.cpp
@@ -0,0 +1,9 @@
+#include <bits/stdc++.h>
+
+using namespace std;
+
+int main() {
+	int a, b;
+	cin >> a >> b;
+	cout << a + b;
+}
diff --git a/tests/packages/ocen/prog/oceningen.cpp b/tests/packages/ocen/prog/oceningen.cpp
new file mode 100644
index 00000000..180f284d
--- /dev/null
+++ b/tests/packages/ocen/prog/oceningen.cpp
@@ -0,0 +1,15 @@
+#include <bits/stdc++.h>
+
+using namespace std;
+
+int main() {
+    ofstream f("ocen0b.in");
+    f << "0 0\n";
+    f.close();
+    f.open("ocen2ocen.in");
+	f << "1 2\n";
+	f.close();
+	f.open("ocen2a.in");
+	f << "1 1\n";
+	f.close();
+}
diff --git a/tests/util.py b/tests/util.py
index fd2d6d33..22c24ca1 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -135,6 +135,14 @@ def get_large_output_package_path():
     """
     return os.path.join(os.path.dirname(__file__), "packages", "large_output")
 
+
+def get_ocen_package_path():
+    """
+    Get path to package for testing ocen archive creation (/tests/packages/ocen)
+    """
+    return os.path.join(os.path.dirname(__file__), "packages", "ocen")
+
+
 def create_ins(package_path, task_id):
     """
     Create .in files for package.