feat: allow Literal types to be classes themselves

corgy/_corgy.py

@@ -248,7 +248,15 @@ class A(Corgy):
 
     The provided values are passed to the `choices` argument of
     `ArgumentParser.add_argument`. All values must be of the same type, which will be
-    inferred from the type of the first value.
+    inferred from the type of the first value. If the first value has a `__bases__`
+    attribute, the type will be inferred as the first base type, and all other choices
+    must be subclasses of that type::
+
+        class A: ...
+        class A1(A): ...
+        class A2(A): ...
+
+        x: Literal[A1, A2]  # inferred type is A
 
     `Literal` itself can be used as a type, for instance inside a `Sequence`::
 
@@ -395,17 +403,26 @@ def add_args_to_parser(cls, parser: argparse.ArgumentParser, name_prefix: str =
                 hasattr(var_base_type, "__origin__")
                 and var_base_type.__origin__ is Literal
             ):
+                # Determine if the first choice has `__bases__`, in which case
+                # the first base class is the type for the argument.
+                try:
+                    c0_type = var_base_type.__args__[0].__bases__[0]
+                except AttributeError:
+                    c0_type = type(var_base_type.__args__[0])
+                    f_check_type: Callable[[Any, Any], bool] = isinstance
+                else:
+                    f_check_type = issubclass
+
                 # All choices must be of the same type.
                 if any(
-                    type(_a) is not type(var_base_type.__args__[0])
-                    for _a in var_base_type.__args__[1:]
+                    not f_check_type(_a, c0_type) for _a in var_base_type.__args__[1:]
                 ):
                     raise TypeError(
-                        f"choices for `{var_name}` not same type: "
+                        f"choices for `{var_name}` not all of type `{c0_type}`: "
                         f"`{var_base_type.__args__}`"
                     )
                 var_choices = var_base_type.__args__
-                var_base_type = type(var_base_type.__args__[0])
+                var_base_type = c0_type
             else:
                 var_choices = None
 

docs/corgy.md

@@ -143,7 +143,17 @@ x: Literal[0, 1, 2]
 
 The provided values are passed to the `choices` argument of
 `ArgumentParser.add_argument`. All values must be of the same type, which will be
-inferred from the type of the first value.
+inferred from the type of the first value. If the first value has a `__bases__`
+attribute, the type will be inferred as the first base type, and all other choices
+must be subclasses of that type:
+
+```python
+class A: ...
+class A1(A): ...
+class A2(A): ...
+
+x: Literal[A1, A2]  # inferred type is A
+```
 
 `Literal` itself can be used as a type, for instance inside a `Sequence`:
 

tests/test_corgy.py

@@ -266,6 +266,40 @@ class C(Corgy):
         with self.assertRaises(TypeError):
             C.add_args_to_parser(self.parser)
 
+    def test_add_args_uses_base_class_for_choice_type(self):
+        class A:
+            pass
+
+        class A1(A):
+            pass
+
+        class A2(A):
+            pass
+
+        class B:
+            pass
+
+        class B1(B):
+            pass
+
+        class BA1(B, A2):
+            pass
+
+        class C(Corgy):
+            x: Literal[A1, A2, BA1]  # type: ignore
+
+        C.add_args_to_parser(self.parser)
+        self.parser.add_argument.assert_called_once_with(
+            "--x", type=A, required=True, choices=(A1, A2, BA1)
+        )
+
+        with self.assertRaises(TypeError):
+
+            class D(Corgy):
+                x: Literal[A1, A2, B1]  # type: ignore
+
+            D.add_args_to_parser(self.parser)
+
     def test_add_args_handles_user_defined_class_as_type(self):
         class T:
             pass