o Moving all doit modules to use DoitConfig().

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])