|
| 1 | +import os |
| 2 | +import glob |
| 3 | +import shutil |
| 4 | +import tarfile |
| 5 | +import argparse |
| 6 | +import yaml |
| 7 | + |
| 8 | +from sinol_make import util |
| 9 | +from sinol_make.helpers import package_util, parsers |
| 10 | +from sinol_make.commands.gen import gen_util |
| 11 | +from sinol_make.interfaces.BaseCommand import BaseCommand |
| 12 | + |
| 13 | + |
| 14 | +class Command(BaseCommand): |
| 15 | + """ |
| 16 | + Class for "export" command. |
| 17 | + """ |
| 18 | + |
| 19 | + def get_name(self): |
| 20 | + return "export" |
| 21 | + |
| 22 | + def configure_subparser(self, subparser: argparse.ArgumentParser): |
| 23 | + parser = subparser.add_parser( |
| 24 | + self.get_name(), |
| 25 | + help='Create archive for oioioi upload', |
| 26 | + description='Creates archive in the current directory ready to upload to sio2 or szkopul.') |
| 27 | + parsers.add_compilation_arguments(parser) |
| 28 | + |
| 29 | + def get_generated_tests(self): |
| 30 | + """ |
| 31 | + Returns list of generated tests. |
| 32 | + Executes ingen to check what tests are generated. |
| 33 | + """ |
| 34 | + if not gen_util.ingen_exists(self.task_id): |
| 35 | + return [] |
| 36 | + |
| 37 | + working_dir = os.path.join(os.getcwd(), 'cache', 'export', 'tests') |
| 38 | + if os.path.exists(working_dir): |
| 39 | + shutil.rmtree(working_dir) |
| 40 | + os.makedirs(working_dir) |
| 41 | + |
| 42 | + ingen_path = gen_util.get_ingen(self.task_id) |
| 43 | + ingen_exe = gen_util.compile_ingen(ingen_path, self.args, self.args.weak_compilation_flags) |
| 44 | + if not gen_util.run_ingen(ingen_exe, working_dir): |
| 45 | + util.exit_with_error('Failed to run ingen.') |
| 46 | + |
| 47 | + tests = glob.glob(os.path.join(working_dir, f'{self.task_id}*.in')) |
| 48 | + return [package_util.extract_test_id(test) for test in tests] |
| 49 | + |
| 50 | + def copy_package_required_files(self, target_dir: str): |
| 51 | + """ |
| 52 | + Copies package files and directories from |
| 53 | + current directory to target directory. |
| 54 | + :param target_dir: Directory to copy files to. |
| 55 | + """ |
| 56 | + files = ['config.yml', 'makefile.in', 'Makefile.in', |
| 57 | + 'prog', 'doc', 'attachments', 'dlazaw'] |
| 58 | + for file in files: |
| 59 | + file_path = os.path.join(os.getcwd(), file) |
| 60 | + if os.path.exists(file_path): |
| 61 | + if os.path.isdir(file_path): |
| 62 | + shutil.copytree(file_path, os.path.join(target_dir, file)) |
| 63 | + else: |
| 64 | + shutil.copy(file_path, target_dir) |
| 65 | + |
| 66 | + print('Copying example tests...') |
| 67 | + for ext in ['in', 'out']: |
| 68 | + os.mkdir(os.path.join(target_dir, ext)) |
| 69 | + for test in glob.glob(os.path.join(os.getcwd(), ext, f'{self.task_id}0*.{ext}')): |
| 70 | + shutil.copy(test, os.path.join(target_dir, ext)) |
| 71 | + |
| 72 | + print('Generating tests...') |
| 73 | + generated_tests = self.get_generated_tests() |
| 74 | + tests_to_copy = [] |
| 75 | + for ext in ['in', 'out']: |
| 76 | + for test in glob.glob(os.path.join(os.getcwd(), ext, f'{self.task_id}*.{ext}')): |
| 77 | + if package_util.extract_test_id(test) not in generated_tests: |
| 78 | + tests_to_copy.append(test) |
| 79 | + |
| 80 | + if len(tests_to_copy) > 0: |
| 81 | + print(util.warning(f'Found {len(tests_to_copy)} tests that are not generated by ingen.')) |
| 82 | + for test in tests_to_copy: |
| 83 | + print(util.warning(f'Coping {os.path.basename(test)}...')) |
| 84 | + shutil.copy(test, os.path.join(target_dir, os.path.splitext(os.path.basename(test))[1])) |
| 85 | + |
| 86 | + def create_makefile_in(self, target_dir: str, config: dict): |
| 87 | + """ |
| 88 | + Creates required `makefile.in` file. |
| 89 | + :param target_dir: Directory to create files in. |
| 90 | + :param config: Config dictionary. |
| 91 | + """ |
| 92 | + with open(os.path.join(target_dir, 'makefile.in'), 'w') as f: |
| 93 | + cxx_flags = '-std=c++17' |
| 94 | + c_flags = '-std=c17' |
| 95 | + if 'extra_compilation_args' in config: |
| 96 | + if 'cpp' in config['extra_compilation_args']: |
| 97 | + cxx_flags += ' ' + ' '.join(config['extra_compilation_args']['cpp']) |
| 98 | + if 'c' in config['extra_compilation_args']: |
| 99 | + c_flags += ' ' + ' '.join(config['extra_compilation_args']['c']) |
| 100 | + |
| 101 | + f.write(f'MODE = wer\n' |
| 102 | + f'ID = {self.task_id}\n' |
| 103 | + f'SIG = sinolmake\n' |
| 104 | + f'\n' |
| 105 | + f'TIMELIMIT = {config["time_limit"]}\n' |
| 106 | + f'SLOW_TIMELIMIT = {4 * config["time_limit"]}\n' |
| 107 | + f'MEMLIMIT = {config["memory_limit"]}\n' |
| 108 | + f'\n' |
| 109 | + f'OI_TIME = oiejq\n' |
| 110 | + f'\n' |
| 111 | + f'CXXFLAGS += {cxx_flags}\n' |
| 112 | + f'CFLAGS += {c_flags}\n') |
| 113 | + |
| 114 | + def compress(self, target_dir): |
| 115 | + """ |
| 116 | + Compresses target directory to archive. |
| 117 | + :param target_dir: Target directory path. |
| 118 | + :return: Path to archive. |
| 119 | + """ |
| 120 | + archive = os.path.join(os.getcwd(), f'{self.task_id}.tgz') |
| 121 | + with tarfile.open(archive, "w:gz") as tar: |
| 122 | + tar.add(target_dir, arcname=os.path.basename(target_dir)) |
| 123 | + return archive |
| 124 | + |
| 125 | + def run(self, args: argparse.Namespace): |
| 126 | + util.exit_if_not_package() |
| 127 | + |
| 128 | + self.args = args |
| 129 | + self.task_id = package_util.get_task_id() |
| 130 | + |
| 131 | + with open(os.path.join(os.getcwd(), 'config.yml'), 'r') as config_file: |
| 132 | + config = yaml.load(config_file, Loader=yaml.FullLoader) |
| 133 | + |
| 134 | + export_package_path = os.path.join(os.getcwd(), 'cache', 'export', self.task_id) |
| 135 | + if os.path.exists(export_package_path): |
| 136 | + shutil.rmtree(export_package_path) |
| 137 | + os.makedirs(export_package_path) |
| 138 | + |
| 139 | + self.copy_package_required_files(export_package_path) |
| 140 | + self.create_makefile_in(export_package_path, config) |
| 141 | + archive = self.compress(export_package_path) |
| 142 | + |
| 143 | + print(util.info(f'Exported to {self.task_id}.tgz')) |
0 commit comments