Fixing a bug in tuples with Literals

swansonk14 · Sep 22, 2023 · 18b13d1 · 18b13d1
1 parent e68e014
commit 18b13d1
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 2 deletions.
diff --git a/tap/tap.py b/tap/tap.py
@@ -239,6 +239,9 @@ def _add_argument(self, *name_or_flags, **kwargs) -> None:
                     else:
                         kwargs['nargs'] = len(types)
 
+                    # Handle Literal types
+                    types = [get_literals(tp, variable)[0] if is_literal_type(tp) else tp for tp in types]
+
                     var_type = TupleTypeEnforcer(types=types, loop=loop)
 
                 if get_origin(var_type) in BOXED_TYPES:

diff --git a/tap/utils.py b/tap/utils.py
@@ -283,7 +283,10 @@ def get_literals(literal: Literal, variable: str) -> Tuple[Callable[[str], Any],
         raise ArgumentTypeError('All literals must have unique string representations')
 
     def var_type(arg: str) -> Any:
-        return str_to_literal.get(arg, arg)
+        if arg not in str_to_literal:
+            raise ArgumentTypeError(f'Value for variable "{variable}" must be one of {literals}.')
+
+        return str_to_literal[arg]
 
     return var_type, literals
 

diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -1170,12 +1170,60 @@ class TupleClassTap(Tap):
         input_args = ('1', '2', '3', '4')
         true_args = (1, '2', Dummy('3'), Dummy('4'))
 
-        args = TupleClassTap().parse_args([
+        args = TupleClassTap().parse_args(['--tup', *input_args])
+
+        self.assertEqual(args.tup, true_args)
+
+    def test_tuple_literally_one(self):
+        class LiterallyOne(Tap):
+            tup: Tuple[Literal[1]]
+
+        input_args = ('1',)
+        true_args = (1,)
+
+        args = LiterallyOne().parse_args(['--tup', *input_args])
+
+        self.assertEqual(args.tup, true_args)
+
+        with self.assertRaises(SystemExit):
+            LiterallyOne().parse_args(['--tup', '2'])
+
+    def test_tuple_literally_two(self):
+        class LiterallyTwo(Tap):
+            tup: Tuple[Literal['two', 2], Literal[2, 'too']]
+
+        input_args = ('2', 'too')
+        true_args = (2, 'too')
+
+        args = LiterallyTwo().parse_args([
             '--tup', *input_args
         ])
 
         self.assertEqual(args.tup, true_args)
 
+        with self.assertRaises(SystemExit):
+            LiterallyTwo().parse_args(['--tup', '2'])
+
+        with self.assertRaises(SystemExit):
+            LiterallyTwo().parse_args(['--tup', '2', '3'])
+
+        with self.assertRaises(SystemExit):
+            LiterallyTwo().parse_args(['--tup', 'too', 'two'])
+
+    def test_tuple_literally_infinity(self):
+        class LiterallyInfinity(Tap):
+            tup: Tuple[Literal['infinity', '∞'], ...]
+
+        input_args = ('∞', 'infinity', '∞')
+        true_args = input_args
+
+        args = LiterallyInfinity().parse_args(['--tup', *input_args])
+
+        self.assertEqual(args.tup, true_args)
+
+        with self.assertRaises(SystemExit):
+            LiterallyInfinity().parse_args(['--tup', '8'])
+
     def test_tuple_wrong_type_fails(self):
         class TupleTapTypeFails(Tap):
             tup: Tuple[int]