Merge pull request #34 from mpkocher/add-enum

Version 3.4.0. Improve support for simple Enums.

CHANGELOG.md

@@ -1,5 +1,9 @@
 # CHANGELOG
 
+## Version 3.4.0
+
+- Improve support for simple `Enum`s. 
+
 ## Version 3.3.0
 
 - Add support for `List` and `Set` fields by [Marius van Niekerk](https://github.com/mariusvniekerk)

pydantic_cli/__init__.py

@@ -3,6 +3,7 @@
 import traceback
 import logging
 import typing as T
+from enum import Enum
 from typing import Callable as F
 import pydantic.fields
 
@@ -227,6 +228,10 @@ def _add_pydantic_field_to_parser(
     else:
         shape_kwargs = {}
 
+    choices: T.Optional[T.List[T.Any]] = None
+    if issubclass(field.type_, Enum):
+        choices = [x.value for x in field.type_.__members__.values()]
+
     if field.type_ == bool:
         __add_boolean_arg_to_parser(
             parser, field_id, cli_custom, default_value, is_required
@@ -240,6 +245,7 @@ def _add_pydantic_field_to_parser(
             default=default_value,
             dest=field_id,
             required=is_required,
+            choices=choices,  # type: ignore
             **shape_kwargs,  # type: ignore
         )
 

pydantic_cli/_version.py

@@ -1 +1 @@
-__version__ = "3.3.0"
+__version__ = "3.4.0"

pydantic_cli/examples/simple_with_enum.py

@@ -0,0 +1,37 @@
+from enum import Enum, auto
+from typing import Set
+from pydantic import BaseModel
+
+from pydantic_cli import run_and_exit
+
+
+class Mode(str, Enum):
+    """Note that if you use `auto`, you must specify the enum
+    using the int value (as a string).
+    """
+
+    alpha = auto()
+    beta = auto()
+
+
+class State(str, Enum):
+    """Note, this is case sensitive when providing it from the commandline"""
+
+    RUNNING = "RUNNING"
+    FAILED = "FAILED"
+    SUCCESSFUL = "SUCCESSFUL"
+
+
+class Options(BaseModel):
+    states: Set[State]
+    mode: Mode
+    max_records: int = 100
+
+
+def example_runner(opts: Options) -> int:
+    print(f"Mock example running with {opts}")
+    return 0
+
+
+if __name__ == "__main__":
+    run_and_exit(Options, example_runner, description=__doc__, version="0.1.0")

pydantic_cli/tests/test_examples_simple_with_custom_and_setup_log.py

@@ -26,5 +26,5 @@ def test_simple_02(self):
     def test_simple_03(self):
         self.run_config(
             ["-i", "/path/to/file.txt", "-m", "1234", "--log_level", "BAD_LOG_LEVEL"],
-            exit_code=1,
+            exit_code=2,
         )

pydantic_cli/tests/test_examples_simple_with_enum.py

@@ -0,0 +1,24 @@
+from . import _TestHarness, HarnessConfig
+
+from pydantic_cli.examples.simple_with_enum import Options, example_runner
+
+
+class TestExamples(_TestHarness[Options]):
+
+    CONFIG = HarnessConfig(Options, example_runner)
+
+    def test_simple_01(self):
+        args = ["--states", "RUNNING", "FAILED", "--max_records", "1234", "--mode", "1"]
+        self.run_config(args)
+
+    def test_bad_enum_value(self):
+        args = [
+            "--states",
+            "RUNNING",
+            "BAD_STATE",
+            "--max_records",
+            "1234",
+            "--mode",
+            "1",
+        ]
+        self.run_config(args, exit_code=2)