ChoiceParameter checks if option is valid within .normalize (#2454)

spotify · Jul 10, 2018 · bdfcb6c · bdfcb6c
1 parent d994bca
commit bdfcb6c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/luigi/parameter.py b/luigi/parameter.py
@@ -1201,8 +1201,11 @@ def __init__(self, var_type=str, *args, **kwargs):
 
     def parse(self, s):
         var = self._var_type(s)
+        return self.normalize(var)
+
+    def normalize(self, var):
         if var in self._choices:
             return var
         else:
-            raise ValueError("{s} is not a valid choice from {choices}".format(
-                s=s, choices=self._choices))
+            raise ValueError("{var} is not a valid choice from {choices}".format(
+                var=var, choices=self._choices))
diff --git a/test/choice_parameter_test.py b/test/choice_parameter_test.py
@@ -53,3 +53,8 @@ def test_serialize_parse(self):
         a = luigi.ChoiceParameter(var_type=str, choices=["1", "2", "3"])
         b = "3"
         self.assertEqual(b, a.parse(a.serialize(b)))
+
+    def test_invalid_choice_task(self):
+        class Foo(luigi.Task):
+            args = luigi.ChoiceParameter(var_type=str, choices=["1", "2", "3"])
+        self.assertRaises(ValueError, lambda: Foo(args="4"))