Training website prototype problem creation script

nztrain.py

#!/usr/bin/env python3

# Problem creation auxiliary script.
# Run: chmod +x nztrain.py && ./nztrain.py
# Configure: model/model.py
#
# Do not move this script as it uses
# paths relative to its file location.

import importlib
import argparse
import shutil
import os
import json

from zipfile import ZipFile
from pathlib import Path

TEMPLATE_MODEL = """\
# Problem model configuration script.
# This script is read and executed for packaging.
#
# If you wish to run additional scripts before packaging
# you may add logic for executing them to this script:
#
#   import os
#   os.system('make')
#
# It is recommended that you compile artifacts to 
# a 'build' folder as any directory with this name 
# is excluded from the final package.
#
# Test cases defined in 'SAMPLES' and 'SETS' must match
# the file names in the 'inputs' and 'outputs' directory
# without the '.txt' extension.

# Sample test cases.
SAMPLES = ['case0']

# Names of prerequisite test case sets.
PREREQUISITES = ['set']

# Test set definitions.
SETS = {
    'set': {
        'points': 100,
        'test_cases': [f'case{i}' for i in range(0, 10)],
    },
}
"""

EXCLUDE = [
    'build',
    '__pycache__',
]


# TODO: add generate stub to model
# TODO: add verify stub to model

def action_new(args):
    root = Path(args.name)
    os.mkdir(root)

    os.mkdir(root / 'inputs')
    os.mkdir(root / 'outputs')
    os.mkdir(root / 'model')

    shutil.copy(__file__, root)
    model_path = root / 'model' / 'model.py'
    with open(model_path, 'w') as f:
        f.write(TEMPLATE_MODEL)

    print(f"Problem '{root}' generated.")
    print(f"Next steps:")
    print(f"1. Configure: '{model_path}'")
    print(f"2. Run: '{root / 'nztrain.py'} package'")


def action_package(_):
    root = Path(__file__).parent
    os.chdir(root / 'model')
    model = importlib.import_module(f'model.model')

    # Write specification file.
    with open(root / 'specification.yaml', 'w') as f:
        f.write("# AUTOGENERATED FILE: see 'nztrain.py'\n")
        json.dump({
            'samples': model.SAMPLES,
            'prerequisites': model.PREREQUISITES,
            'test_sets': model.SETS,
        }, f, indent=4)

    # Package problem contents.
    build = root / 'build'
    package_path = build / 'package.zip'
    build.mkdir(parents=True, exist_ok=True)
    with ZipFile(package_path, 'w') as z:
        for base, folders, files in os.walk(root):
            files = [f for f in files if not f[0] == '.']
            folders[:] = [f for f in folders if not f[0] == '.']
            folders[:] = [f for f in folders if f not in EXCLUDE]

            for name in folders + files:
                path = Path(base) / name
                relative = path.relative_to(root)
                z.write(path, relative)

    print(f"Package generated: '{package_path}'")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Training website problem creation tool.')
    parser.set_defaults(action=lambda _: parser.print_help())
    parsers = parser.add_subparsers(help="command help")

    package = parsers.add_parser('package', help='compile the problem into a package')
    package.set_defaults(action=action_package)

    new = parsers.add_parser('new', help='generate the template for a new problem')
    new.add_argument('name', help='name of generated problem')
    new.set_defaults(action=action_new)

    arguments = parser.parse_args()
    arguments.action(arguments)