Support enums

Author: timkpaine

hatch_build/cli.py

@@ -1,4 +1,5 @@
 from argparse import ArgumentParser
+from enum import Enum
 from logging import Formatter, StreamHandler, getLogger
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Dict, List, Literal, Optional, Tuple, Type, Union, get_args, get_origin
@@ -98,6 +99,11 @@ def _recurse_add_fields(parser: ArgumentParser, model: Union["BaseModel", Type["
             except TypeError:
                 # TODO: handle more complex types if needed
                 parser.add_argument(arg_name, type=str, default=default_value)
+        elif isinstance(field_type, type) and issubclass(field_type, Enum):
+            #############
+            # MARK: Enum
+            enum_choices = [e.value for e in field_type]
+            parser.add_argument(arg_name, type=type(enum_choices[0]), choices=enum_choices, default=default_value)
         elif isinstance(field_type, type) and issubclass(field_type, Path):
             #############
             # MARK: Path

hatch_build/tests/test_cli_model.py

@@ -1,4 +1,5 @@
 import sys
+from enum import Enum
 from pathlib import Path
 from typing import Dict, List, Literal, Optional
 from unittest.mock import patch
@@ -8,6 +9,11 @@
 from hatch_build.cli import hatchling, parse_extra_args_model
 
 
+class MyEnum(Enum):
+    OPTION_A = "option_a"
+    OPTION_B = "option_b"
+
+
 class SubModel(BaseModel, validate_assignment=True):
     sub_arg: int = 42
     sub_arg_with_value: str = "sub_default"
@@ -19,6 +25,7 @@ class MyTopLevelModel(BaseModel, validate_assignment=True):
     extra_arg_with_value_equals: Optional[str] = "default_equals"
     extra_arg_literal: Literal["a", "b", "c"] = "a"
 
+    enum_arg: MyEnum = MyEnum.OPTION_A
     list_arg: List[int] = [1, 2, 3]
     dict_arg: Dict[str, str] = {}
     dict_arg_default_values: Dict[str, str] = {"existing-key": "existing-value"}
@@ -33,6 +40,8 @@ class MyTopLevelModel(BaseModel, validate_assignment=True):
     submodel_dict: Dict[str, SubModel] = {}
     submodel_dict_instanced: Dict[str, SubModel] = {"a": SubModel()}
 
+    unsupported_literal: Literal[b"test"] = b"test"
+
 
 class TestCLIMdel:
     def test_get_arg_from_model(self):
@@ -49,6 +58,8 @@ def test_get_arg_from_model(self):
                 "--extra-arg-not-in-parser",
                 "--extra-arg-literal",
                 "b",
+                "--enum-arg",
+                "option_b",
                 "--list-arg",
                 "1,2,3",
                 "--dict-arg",
@@ -80,6 +91,7 @@ def test_get_arg_from_model(self):
         assert model.extra_arg_with_value == "value"
         assert model.extra_arg_with_value_equals == "value2"
         assert model.extra_arg_literal == "b"
+        assert model.enum_arg == MyEnum.OPTION_B
         assert model.list_arg == [1, 2, 3]
         assert model.dict_arg == {"key1": "value1", "key2": "value2"}
         assert model.dict_arg_default_values == {"existing-key": "new-value"}