Feature/ Runtime validation of TaskInstanceParameter() and ListTaskInstanceParameter() by subclass bound

gokart/parameter.py

@@ -5,11 +5,21 @@
 import luigi
 from luigi import task_register
 
+import gokart
+
 logger = getLogger(__name__)
 
 
 class TaskInstanceParameter(luigi.Parameter):
 
+    def __init__(self, *args, **kwargs):
+        bound = kwargs.pop('bound', gokart.TaskOnKart)
+        if isinstance(bound, type):
+            self._bound = bound
+        else:
+            raise ValueError(f'bound must be a type, not {type(bound)}')
+        super().__init__(*args, **kwargs)
+
     @staticmethod
     def _recursive(param_dict):
         params = param_dict['params']
@@ -36,6 +46,11 @@ def serialize(self, x):
         values = dict(type=x.get_task_family(), params=params)
         return luigi.DictParameter().serialize(values)
 
+    def normalize(self, v):
+        if not isinstance(v, self._bound):
+            raise ValueError(f'{v} is not an instance of {self._bound}')
+        return v
+
 
 class _TaskInstanceEncoder(json.JSONEncoder):
 
@@ -48,12 +63,26 @@ def default(self, obj):
 
 class ListTaskInstanceParameter(luigi.Parameter):
 
+    def __init__(self, *args, **kwargs):
+        bound = kwargs.pop('bound', gokart.TaskOnKart)
+        if isinstance(bound, type):
+            self._bound = bound
+        else:
+            raise ValueError(f'bound must be a type, not {type(bound)}')
+        super().__init__(*args, **kwargs)
+
     def parse(self, s):
         return [TaskInstanceParameter().parse(x) for x in list(json.loads(s))]
 
     def serialize(self, x):
         return json.dumps(x, cls=_TaskInstanceEncoder)
 
+    def normalize(self, values):
+        for v in values:
+            if not isinstance(v, self._bound):
+                raise ValueError(f'{v} is not an instance of {self._bound}')
+        return values
+
 
 class ExplicitBoolParameter(luigi.BoolParameter):
 

test/test_task_instance_parameter.py

@@ -11,6 +11,16 @@ class _DummySubTask(TaskOnKart):
     pass
 
 
+class _DummyCorrectSubClassTask(_DummySubTask):
+    task_namespace = __name__
+    pass
+
+
+class _DummyInvalidSubClassTask(TaskOnKart):
+    task_namespace = __name__
+    pass
+
+
 class _DummyTask(TaskOnKart):
     task_namespace = __name__
     param = luigi.IntParameter()
@@ -40,6 +50,54 @@ def test_serialize_and_parse_list_params(self):
         parsed = gokart.TaskInstanceParameter().parse(s)
         self.assertEqual(parsed.task_id, original.task_id)
 
+    def test_invalid_bound(self):
+        self.assertRaises(ValueError, lambda: gokart.TaskInstanceParameter(bound=1))  # not type instance
+
+    def test_params_with_correct_subclass_bound(self):
+
+        class _DummyPipelineA(TaskOnKart):
+            task_namespace = __name__
+            subtask = gokart.TaskInstanceParameter(bound=_DummySubTask)
+
+        task = _DummyPipelineA(subtask=_DummyCorrectSubClassTask())
+        self.assertEqual(task.requires()['subtask'], _DummyCorrectSubClassTask())
+
+    def test_params_with_invalid_subclass_bound(self):
+
+        class _DummyPipelineB(TaskOnKart):
+            task_namespace = __name__
+            subtask = gokart.TaskInstanceParameter(bound=_DummySubTask)
+
+        with self.assertRaises(ValueError):
+            _DummyPipelineB(subtask=_DummyInvalidSubClassTask())
+
+
+class ListTaskInstanceParameterTest(unittest.TestCase):
+
+    def setUp(self):
+        _DummyTask.clear_instance_cache()
+
+    def test_invalid_bound(self):
+        self.assertRaises(ValueError, lambda: gokart.ListTaskInstanceParameter(bound=1))  # not type instance
+
+    def test_list_params_with_correct_subclass_bound(self):
+
+        class _DummyPipelineC(TaskOnKart):
+            task_namespace = __name__
+            subtask = gokart.ListTaskInstanceParameter(bound=_DummySubTask)
+
+        task = _DummyPipelineC(subtask=[_DummyCorrectSubClassTask()])
+        self.assertEqual(task.requires()['subtask'], (_DummyCorrectSubClassTask(), ))
+
+    def test_list_params_with_invalid_subclass_bound(self):
+
+        class _DummyPipelineD(TaskOnKart):
+            task_namespace = __name__
+            subtask = gokart.ListTaskInstanceParameter(bound=_DummySubTask)
+
+        with self.assertRaises(ValueError):
+            _DummyPipelineD(subtask=[_DummyInvalidSubClassTask(), _DummyCorrectSubClassTask])
+
 
 if __name__ == '__main__':
     unittest.main()