Adding Union type tests, raising an error when a type function is not provided for union types, and fixing exception types in tests

Author: swansonk14

.github/workflows/code-coverage.yml

@@ -6,13 +6,13 @@ jobs:
   run:
     runs-on: ubuntu-latest
     env:
-      PYTHON: '3.9'
+      PYTHON: '3.10'
     steps:
     - uses: actions/checkout@main
     - name: Setup Python
       uses: actions/setup-python@main
       with:
-        python-version: 3.9
+        python-version: 3.10
     - name: Generate coverage report
       run: |
         git config --global user.email "[email protected]"

.github/workflows/tests.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: [3.6, 3.7, 3.8, 3.9, 3.10]
 
     steps:
     - uses: actions/checkout@v2

tap/tap.py

@@ -10,11 +10,12 @@
 import time
 from types import MethodType
 from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, TypeVar, Union, get_type_hints
-from typing_inspect import is_literal_type, get_args
+from typing_inspect import is_literal_type
 from warnings import warn
 
 from tap.utils import (
     get_class_variables,
+    get_args,
     get_argument_name,
     get_dest,
     get_origin,
@@ -37,7 +38,8 @@
 # Constants
 EMPTY_TYPE = get_args(List)[0] if len(get_args(List)) > 0 else tuple()
 BOXED_COLLECTION_TYPES = {List, list, Set, set, Tuple, tuple}
-OPTIONAL_TYPES = {Optional, Union} | ({UnionType} if sys.version_info >= (3, 10) else set())
+UNION_TYPES = {Union} | ({UnionType} if sys.version_info >= (3, 10) else set())
+OPTIONAL_TYPES = {Optional} | UNION_TYPES
 BOXED_TYPES = BOXED_COLLECTION_TYPES | OPTIONAL_TYPES
 
 
@@ -172,11 +174,26 @@ def _add_argument(self, *name_or_flags, **kwargs) -> None:
 
             # If type is not explicitly provided, set it if it's one of our supported default types
             if 'type' not in kwargs:
-
                 # Unbox Union[type] (Optional[type]) and set var_type = type
                 if get_origin(var_type) in OPTIONAL_TYPES:
                     var_args = get_args(var_type)
 
+                    # If type is Union or Optional without inner types, set type to equivalent of Optional[str]
+                    if len(var_args) == 0:
+                        var_args = (str, type(None))
+
+                    # Raise error if type function is not explicitly provided for Union types (not including Optionals)
+                    if get_origin(var_type) in UNION_TYPES and not (len(var_args) == 2 and var_args[1] == type(None)):
+                        raise ArgumentTypeError(
+                            'For Union types, you must include an explicit type function in the configure method. '
+                            'For example,\n\n'
+                            'class Args(Tap):\n'
+                            '    arg: Union[int, float]\n'
+                            '\n'
+                            '    def configure(self) -> None:\n'
+                            '        self.add_argument("--arg", type=lambda x: float(x) if "." in x else int(x))'
+                        )
+
                     if len(var_args) > 0:
                         var_type = var_args[0]
 
@@ -242,7 +259,7 @@ def _add_argument(self, *name_or_flags, **kwargs) -> None:
                         kwargs['type'] = boolean_type
                         kwargs['choices'] = [True, False]  # this makes the help message more helpful
                     else:
-                        action_cond = "true" if kwargs.get("required", False) or not kwargs["default"] else "false"
+                        action_cond = 'true' if kwargs.get('required', False) or not kwargs['default'] else 'false'
                         kwargs['action'] = kwargs.get('action', f'store_{action_cond}')
                 elif kwargs.get('action') not in {'count', 'append_const'}:
                     kwargs['type'] = var_type

tap/utils.py

@@ -24,7 +24,7 @@
     Union,
 )
 from typing_extensions import Literal
-from typing_inspect import get_args, get_origin as typing_inspect_get_origin
+from typing_inspect import get_args as typing_inspect_get_args, get_origin as typing_inspect_get_origin
 
 if sys.version_info >= (3, 10):
     from types import UnionType
@@ -487,3 +487,12 @@ def get_origin(tp: Any) -> Any:
         origin = UnionType
 
     return origin
+
+
+# TODO: remove this once typing_insepct.get_args is fixed for Python 3.10 union types
+def get_args(tp: Any) -> Tuple[type, ...]:
+    """Same as typing_inspect.get_args but fixes Python 3.10 union types."""
+    if sys.version_info >= (3, 10) and isinstance(tp, UnionType):
+        return tp.__args__
+
+    return typing_inspect_get_args(tp)

tests/test_integration.py

@@ -1,10 +1,11 @@
+from argparse import ArgumentTypeError
 from copy import deepcopy
 import os
 from pathlib import Path
 import pickle
 import sys
 from tempfile import TemporaryDirectory
-from typing import Any, Iterable, List, Optional, Set, Tuple
+from typing import Any, Iterable, List, Optional, Set, Tuple, Union
 from typing_extensions import Literal
 import unittest
 from unittest import TestCase
@@ -556,6 +557,186 @@ def write(self, msg):
             LiteralCrashTap().parse_args(['--arg_lit', '123'])
 
 
+def convert_str_or_int(str_or_int: str) -> Union[str, int]:
+    try:
+        return int(str_or_int)
+    except ValueError:
+        return str_or_int
+
+
+def convert_person_or_str(person_or_str: str) -> Union[Person, str]:
+    if person_or_str == person_or_str.title():
+        return Person(person_or_str)
+
+    return person_or_str
+
+
+def convert_many_types(input_str: str) -> Union[int, float, Person, str]:
+    try:
+        return int(input_str)
+    except ValueError:
+        try:
+            return float(input_str)
+        except ValueError:
+            if input_str == input_str.title():
+                return Person(input_str)
+
+            return input_str
+
+
+# TODO: test crash if not specifying type function
+class UnionTypeTap(Tap):
+    union_zero_required_arg: Union
+    union_zero_default_arg: Union = 'hi'
+    union_one_required_arg: Union[str]
+    union_one_default_arg: Union[str] = 'there'
+    union_two_required_arg: Union[str, int]
+    union_two_default_int_arg: Union[str, int] = 5
+    union_two_default_str_arg: Union[str, int] = 'year old'
+    union_custom_required_arg: Union[Person, str]
+    union_custom_required_flip_arg: Union[str, Person]
+    union_custom_default_arg: Union[Person, str] = Person('Jesse')
+    union_custom_default_flip_arg: Union[str, Person] = 'I want'
+    union_none_required_arg: Union[int, None]
+    union_none_required_flip_arg: Union[None, int]
+    union_many_required_arg: Union[int, float, Person, str]
+    union_many_default_arg: Union[int, float, Person, str] = 3.14
+
+    def configure(self) -> None:
+        self.add_argument('--union_two_required_arg', type=convert_str_or_int)
+        self.add_argument('--union_two_default_int_arg', type=convert_str_or_int)
+        self.add_argument('--union_two_default_str_arg', type=convert_str_or_int)
+        self.add_argument('--union_custom_required_arg', type=convert_person_or_str)
+        self.add_argument('--union_custom_required_flip_arg', type=convert_person_or_str)
+        self.add_argument('--union_custom_default_arg', type=convert_person_or_str)
+        self.add_argument('--union_custom_default_flip_arg', type=convert_person_or_str)
+        self.add_argument('--union_none_required_flip_arg', type=int)
+        self.add_argument('--union_many_required_arg', type=convert_many_types)
+        self.add_argument('--union_many_default_arg', type=convert_many_types)
+
+
+if sys.version_info >= (3, 10):
+    class UnionType310Tap(Tap):
+        union_two_required_arg: str | int
+        union_two_default_int_arg: str | int = 10
+        union_two_default_str_arg: str | int = 'pieces of pie for'
+        union_custom_required_arg: Person | str
+        union_custom_required_flip_arg: str | Person
+        union_custom_default_arg: Person | str = Person('Kyle')
+        union_custom_default_flip_arg: str | Person = 'making'
+        union_none_required_arg: int | None
+        union_none_required_flip_arg: None | int
+        union_many_required_arg: int | float | Person | str
+        union_many_default_arg: int | float | Person | str = 3.14 * 10 / 8
+
+        def configure(self) -> None:
+            self.add_argument('--union_two_required_arg', type=convert_str_or_int)
+            self.add_argument('--union_two_default_int_arg', type=convert_str_or_int)
+            self.add_argument('--union_two_default_str_arg', type=convert_str_or_int)
+            self.add_argument('--union_custom_required_arg', type=convert_person_or_str)
+            self.add_argument('--union_custom_required_flip_arg', type=convert_person_or_str)
+            self.add_argument('--union_custom_default_arg', type=convert_person_or_str)
+            self.add_argument('--union_custom_default_flip_arg', type=convert_person_or_str)
+            self.add_argument('--union_none_required_flip_arg', type=int)
+            self.add_argument('--union_many_required_arg', type=convert_many_types)
+            self.add_argument('--union_many_default_arg', type=convert_many_types)
+
+
+class UnionTypeTests(TestCase):
+    def test_union_types(self):
+        union_zero_required_arg = 'Kyle'
+        union_one_required_arg = 'ate'
+        union_two_required_arg = '2'
+        union_custom_required_arg = 'many'
+        union_custom_required_flip_arg = 'Jesse'
+        union_none_required_arg = '1'
+        union_none_required_flip_arg = '5'
+        union_many_required_arg = 'still hungry'
+
+        args = UnionTypeTap().parse_args([
+            '--union_zero_required_arg', union_zero_required_arg,
+            '--union_one_required_arg', union_one_required_arg,
+            '--union_two_required_arg', union_two_required_arg,
+            '--union_custom_required_arg', union_custom_required_arg,
+            '--union_custom_required_flip_arg', union_custom_required_flip_arg,
+            '--union_none_required_arg', union_none_required_arg,
+            '--union_none_required_flip_arg', union_none_required_flip_arg,
+            '--union_many_required_arg', union_many_required_arg
+        ])
+
+        union_two_required_arg = int(union_two_required_arg)
+        union_custom_required_flip_arg = Person(union_custom_required_flip_arg)
+        union_none_required_arg = int(union_none_required_arg)
+        union_none_required_flip_arg = int(union_none_required_flip_arg)
+
+        self.assertEqual(args.union_zero_required_arg, union_zero_required_arg)
+        self.assertEqual(args.union_zero_default_arg, UnionTypeTap.union_zero_default_arg)
+        self.assertEqual(args.union_one_required_arg, union_one_required_arg)
+        self.assertEqual(args.union_one_default_arg, UnionTypeTap.union_one_default_arg)
+        self.assertEqual(args.union_two_required_arg, union_two_required_arg)
+        self.assertEqual(args.union_two_default_int_arg, UnionTypeTap.union_two_default_int_arg)
+        self.assertEqual(args.union_two_default_str_arg, UnionTypeTap.union_two_default_str_arg)
+        self.assertEqual(args.union_custom_required_arg, union_custom_required_arg)
+        self.assertEqual(args.union_custom_required_flip_arg, union_custom_required_flip_arg)
+        self.assertEqual(args.union_custom_default_arg, UnionTypeTap.union_custom_default_arg)
+        self.assertEqual(args.union_custom_default_flip_arg, UnionTypeTap.union_custom_default_flip_arg)
+        self.assertEqual(args.union_none_required_arg, union_none_required_arg)
+        self.assertEqual(args.union_none_required_flip_arg, union_none_required_flip_arg)
+        self.assertEqual(args.union_many_required_arg, union_many_required_arg)
+        self.assertEqual(args.union_many_default_arg, UnionTypeTap.union_many_default_arg)
+
+    def test_union_missing_type_function(self):
+        class UnionMissingTypeFunctionTap(Tap):
+            arg: Union[int, float]
+
+        with self.assertRaises(ArgumentTypeError):
+            UnionMissingTypeFunctionTap()
+
+    @unittest.skipIf(sys.version_info < (3, 10), 'Union type operator "|" introduced in Python 3.10')
+    def test_union_types_310(self):
+        union_two_required_arg = '1'  # int
+        union_custom_required_arg = 'hungry'  # str
+        union_custom_required_flip_arg = 'Loser'  # Person
+        union_none_required_arg = '8'  # int
+        union_none_required_flip_arg = '100'  # int
+        union_many_required_arg = '3.14'  # float
+
+        args = UnionType310Tap().parse_args([
+            '--union_two_required_arg', union_two_required_arg,
+            '--union_custom_required_arg', union_custom_required_arg,
+            '--union_custom_required_flip_arg', union_custom_required_flip_arg,
+            '--union_none_required_arg', union_none_required_arg,
+            '--union_none_required_flip_arg', union_none_required_flip_arg,
+            '--union_many_required_arg', union_many_required_arg
+        ])
+
+        union_two_required_arg = int(union_two_required_arg)
+        union_custom_required_flip_arg = Person(union_custom_required_flip_arg)
+        union_none_required_arg = int(union_none_required_arg)
+        union_none_required_flip_arg = int(union_none_required_flip_arg)
+        union_many_required_arg = float(union_many_required_arg)
+
+        self.assertEqual(args.union_two_required_arg, union_two_required_arg)
+        self.assertEqual(args.union_two_default_int_arg, UnionType310Tap.union_two_default_int_arg)
+        self.assertEqual(args.union_two_default_str_arg, UnionType310Tap.union_two_default_str_arg)
+        self.assertEqual(args.union_custom_required_arg, union_custom_required_arg)
+        self.assertEqual(args.union_custom_required_flip_arg, union_custom_required_flip_arg)
+        self.assertEqual(args.union_custom_default_arg, UnionType310Tap.union_custom_default_arg)
+        self.assertEqual(args.union_custom_default_flip_arg, UnionType310Tap.union_custom_default_flip_arg)
+        self.assertEqual(args.union_none_required_arg, union_none_required_arg)
+        self.assertEqual(args.union_none_required_flip_arg, union_none_required_flip_arg)
+        self.assertEqual(args.union_many_required_arg, union_many_required_arg)
+        self.assertEqual(args.union_many_default_arg, UnionType310Tap.union_many_default_arg)
+
+    @unittest.skipIf(sys.version_info < (3, 10), 'Union type operator "|" introduced in Python 3.10')
+    def test_union_missing_type_function_310(self):
+        class UnionMissingTypeFunctionTap(Tap):
+            arg: int | float
+
+        with self.assertRaises(ArgumentTypeError):
+            UnionMissingTypeFunctionTap()
+
+
 class AddArgumentTests(TestCase):
     def setUp(self) -> None:
         # Suppress prints from SystemExit
@@ -1044,7 +1225,7 @@ def test_empty_tuple_fails(self):
         class EmptyTupleTap(Tap):
             tup: Tuple[()]
 
-        with self.assertRaises(ValueError):
+        with self.assertRaises(ArgumentTypeError):
             EmptyTupleTap().parse_args([])
 
     def test_tuple_non_tuple_default(self):
@@ -1380,14 +1561,5 @@ def test_pickle(self):
         self.assertEqual(loaded_args.as_dict(), args.as_dict())
 
 
-"""
-- crash if default type not supported
-- user specifying process_args
-- test get reproducibility info
-- test str?
-- test comments
-"""
-
-
 if __name__ == '__main__':
     unittest.main()

tests/test_utils.py

@@ -1,3 +1,4 @@
+from argparse import ArgumentTypeError
 from collections import OrderedDict
 import json
 import os
@@ -297,7 +298,7 @@ def test_get_literals_primitives(self) -> None:
         self.assertEqual([literal_f(str(p)) for p in prims], literals)
 
     def test_get_literals_uniqueness(self) -> None:
-        with self.assertRaises(ValueError):
+        with self.assertRaises(ArgumentTypeError):
             get_literals(Literal['two', 2, '2'], 'number')
 
     def test_get_literals_empty(self) -> None: