diff options
Diffstat (limited to 'src/ee/doit.py')
| -rw-r--r-- | src/ee/doit.py | 11 |
1 files changed, 9 insertions, 2 deletions
diff --git a/src/ee/doit.py b/src/ee/doit.py index dc89fae..c0a01ea 100644 --- a/src/ee/doit.py +++ b/src/ee/doit.py @@ -33,7 +33,10 @@ class DoItConfig(object): return self._dsm def data_sets_for(self, task): - return self._data_sets[task] + try: + return self._data_sets[task] + except KeyError: + raise KeyError("No such task registered in this module: {}".format(task)) def out_data_set_for(self, task): return self.data_sets_for(task)[0] @@ -42,4 +45,8 @@ class DoItConfig(object): return self.data_sets_for(task)[1] def set_data_sets_for(self, task, out_dataset: str, *in_datasets: str): - self._data_sets[task] = [out_dataset, in_datasets] + self._data_sets[task] = [out_dataset, list(in_datasets)] + + def change_data_sets_for_task(self, task, _callable): + ds = self._data_sets[task] + ds[1] = _callable(ds[1]) |