taskmaker.py 1.51 KB
Newer Older
1 2 3 4 5
import sys
import os
import yaml
import numpy

Lukas Weber's avatar
Lukas Weber committed
6 7 8 9 10
try:
    from yaml import CSafeLoader as SafeLoader, CSafeDumper as SafeDumper
except ImportError:
    from yaml import SafeLoader, SafeDumper

11 12 13
def _expand_path(path):
    return os.path.abspath(os.path.expandvars(os.path.expanduser(path)))

14 15
def JobConfig(filename):
    with open(_expand_path(filename), 'r') as f:
Lukas Weber's avatar
Lukas Weber committed
16
        return yaml.load(f, Loader=SafeLoader)
17

18 19
class TaskMaker:
    def __init__(self, name, jobconfig):
20
        self._tm_name = name
21
        self._tm_tasks = []
22 23 24 25
        if type(jobconfig) == dict:
            self._tm_jobconfig = jobconfig
        elif type(jobconfig) == str:
            self._tm_jobconfig = JobConfig(jobconfig)
26 27 28

        self._tm_jobconfig['mc_binary'] = _expand_path(self._tm_jobconfig['mc_binary'])

29 30 31 32 33
    def task(self, **additional_parameters):
        self._tm_tasks.append({**self.__dict__, **additional_parameters})

    def write(self):
        filenames = []
34
        jobfile_dict = { 'jobname': self._tm_name, 'jobconfig': self._tm_jobconfig, 'tasks': {}}
35

36 37 38 39 40 41 42 43 44 45 46 47 48
        for i, t in enumerate(self._tm_tasks):
            task_name = f'task{i+1:04d}'
            task_dict = {}
            for k, v in t.items():
                if k.startswith('_tm_'):
                    continue

                if isinstance(v, numpy.generic): # cast numpy types to their native forms
                    v = v.item()

                task_dict[k] = v
            jobfile_dict['tasks'][task_name] = task_dict

Lukas Weber's avatar
Lukas Weber committed
49
        print(yaml.dump(jobfile_dict, Dumper=SafeDumper))