Add EnumChoice type

CHANGES.rst

@@ -5,6 +5,10 @@ Version 8.2.0
 
 Unreleased
 
+-   Add ``EnumChoice`` ``ParamType`` that can accept an ``enum.Enum`` type
+    and use its keys or values as choices that map back to the ``Enum``.
+    :issue:`605`
+
 
 Version 8.1.3
 -------------

docs/options.rst

@@ -414,6 +414,85 @@ Choices should be unique after considering the effects of
 
 .. _option-prompting:
 
+.. _enum-choice-opts:
+
+Enum Choice Options
+--------------
+
+Similar to the :class:`Choice` type, the :class:`EnumChoice` type may
+be used to have the parameter be a choice of :class:`enum.Enum` names or values
+that may to `enum.Enum` objects when parsed.  It can be instantiated
+with an `enum.Enum` class.  The enum object corresponding to the passed name
+or value (depending on settings) will be returned.
+
+Example:
+
+.. click:example::
+
+    class MockEnum(Enum):
+        FOO = "foo"
+        BAR = "bar"
+        BAZ = "baz"
+
+    @click.command()
+    @click.option('--foo-opt',
+                  type=click.EnumChoice(MockEnum)
+    def handle_foo(foo_opt):
+        click.echo(foo_opt)
+
+What it looks like:
+
+.. click:run::
+
+    invoke(handle_foo, args=['--foo-opt=FOO'])
+    println()
+    invoke(handle_foo, args=['--foo-opt=BAR'])
+    println()
+    invoke(handle_foo, args=['--foo-opt=foo'])
+    println()
+    invoke(handle_foo, args=['--foo-opt=other'])
+    println()
+    invoke(handle_foo, args=['--help'])
+
+Using the ``use_value`` parameter will cause the parser to use the enum
+values as choices instead of the names. All enum values must be unique. Example:
+
+.. click:example::
+
+    class MockEnum(Enum):
+        FOO = "foo"
+        BAR = "bar"
+        BAZ = "baz"
+
+    @click.command()
+    @click.option('--foo-opt',
+                  type=click.EnumChoice(MockEnum, use_value=True)
+    def handle_foo(foo_opt):
+        click.echo(foo_opt)
+
+What it looks like:
+
+.. click:run::
+
+    invoke(handle_foo, args=['--foo-opt=foo'])
+    println()
+    invoke(handle_foo, args=['--foo-opt=bar'])
+    println()
+    invoke(handle_foo, args=['--foo-opt=FOO'])
+    println()
+    invoke(handle_foo, args=['--foo-opt=other'])
+    println()
+    invoke(handle_foo, args=['--help'])
+
+Choices work with options that have ``multiple=True``. If a ``default``
+value is given with ``multiple=True``, it should be a list or tuple of
+valid choices.
+
+The ``case_sensitive`` parameter works identically to the :class:`Choice` type,
+operating whatever the choices derived from the enum are.
+
+.. versionadded:: 8.2
+
 Prompting
 ---------
 

src/click/__init__.py

@@ -52,6 +52,7 @@
 from .types import BOOL as BOOL
 from .types import Choice as Choice
 from .types import DateTime as DateTime
+from .types import EnumChoice as EnumChoice
 from .types import File as File
 from .types import FLOAT as FLOAT
 from .types import FloatRange as FloatRange

src/click/types.py

@@ -1,3 +1,4 @@
+import enum
 import os
 import stat
 import typing as t
@@ -329,6 +330,64 @@ def shell_complete(
         return [CompletionItem(c) for c in matched]
 
 
+class EnumChoice(Choice):
+    """The :class:`EnumChoice` type allows a value to be checked against either the
+    keys or the values of an :class:`~enum.Enum`.
+
+    If using :param:`use_value` is `True`, all enum values should be unique.
+    The inputted value will then be converted to the corresponding enum object.
+
+    The resulting value will be an instance of the Enum class corresponding
+    to the user's key or value passed in.
+
+    See :ref:`choice-opts` for an example.
+
+    :param choices: `enum.Enum` type that will be used to create CLI choices.
+    :param case_sensitive: Set to `False` to make choices case
+        insensitive. Defaults to `True`.
+    :param use_value: Set to `True` to use enum values as choices instead
+        of enum keys. Defaults to `False`.
+
+    :raises ValueError: Raised when `use_value` is True but enum values are not
+        unique.
+
+    .. versionadded:: 8.2
+    """
+
+    name = "enum_choice"
+
+    def __init__(
+        self,
+        choices: t.Type[enum.Enum],
+        case_sensitive: bool = True,
+        use_value: bool = False,
+    ) -> None:
+        self.use_value: bool = use_value
+        if self.use_value:
+            try:
+                enum.unique(choices)
+            except ValueError as err:
+                raise ValueError(
+                    "All values in `choices` must be unique if `use_value` is `True`"
+                ) from err
+            self.choice_to_enum = {str(enum.value): enum for enum in choices}
+        else:
+            self.choice_to_enum = {str(enum.name): enum for enum in choices}
+        super().__init__(
+            choices=list(self.choice_to_enum),
+            case_sensitive=case_sensitive,
+        )
+
+    def convert(
+        self, value: t.Any, param: t.Optional["Parameter"], ctx: t.Optional["Context"]
+    ) -> t.Any:
+        normed_value = super().convert(value, param, ctx)
+        return self.choice_to_enum.get(normed_value)
+
+    def __repr__(self) -> str:
+        return f"EnumChoice({list(self.choices)})"
+
+
 class DateTime(ParamType):
     """The DateTime type converts date strings into `datetime` objects.
 

tests/test_basic.py

@@ -1,4 +1,5 @@
 import os
+from enum import Enum
 from itertools import chain
 
 import pytest
@@ -403,6 +404,113 @@ def cli(method):
     assert "{foo|bar|baz}" in result.output
 
 
+def test_enum_choice_argument(runner):
+    class MockEnum(Enum):
+        foo = 1
+        bar = 2
+        baz = 3
+
+    @click.command()
+    @click.argument("method", type=click.EnumChoice(MockEnum))
+    def cli(method):
+        click.echo(method)
+
+    result = runner.invoke(cli, ["foo"])
+    assert not result.exception
+    assert result.output == "MockEnum.foo\n"
+
+    result = runner.invoke(cli, ["meh"])
+    assert result.exit_code == 2
+    assert (
+        "Invalid value for '{foo|bar|baz}': 'meh' is not one of 'foo',"
+        " 'bar', 'baz'." in result.output
+    )
+
+    result = runner.invoke(cli, ["--help"])
+    assert "{foo|bar|baz}" in result.output
+
+
+def test_enum_choice_argument_using_values(runner):
+    class MockEnum(Enum):
+        foo = 1
+        bar = 2
+        baz = 3
+
+    @click.command()
+    @click.argument("method", type=click.EnumChoice(MockEnum, use_value=True))
+    def cli(method):
+        click.echo(method)
+
+    result = runner.invoke(cli, ["1"])
+    assert not result.exception
+    assert result.output == "MockEnum.foo\n"
+
+    result = runner.invoke(cli, ["4"])
+    assert result.exit_code == 2
+    assert (
+        "Invalid value for '{1|2|3}': '4' is not one of '1',"
+        " '2', '3'." in result.output
+    )
+
+    result = runner.invoke(cli, ["--help"])
+    assert "{1|2|3}" in result.output
+
+
+def test_enum_choice_argument_case_insensitive(runner):
+    class MockEnum(Enum):
+        foo = 1
+        bar = 2
+        baz = 3
+
+    @click.command()
+    @click.argument("method", type=click.EnumChoice(MockEnum, case_sensitive=False))
+    def cli(method):
+        click.echo(method)
+
+    result = runner.invoke(cli, ["FOO"])
+    assert not result.exception
+    assert result.output == "MockEnum.foo\n"
+
+
+def test_enum_choice_argument_case_insensitive_values(runner):
+    class MockEnum(Enum):
+        foo = "foo"
+        bar = "bar"
+        baz = "baz"
+
+    @click.command()
+    @click.argument("method", type=click.EnumChoice(MockEnum, case_sensitive=False))
+    def cli(method):
+        click.echo(method)
+
+    result = runner.invoke(cli, ["FOO"])
+    assert not result.exception
+    assert result.output == "MockEnum.foo\n"
+
+
+def test_enum_choice_argument_non_unique_values(runner):
+    class MockEnum(Enum):
+        foo = 1
+        bar = 2
+        baz = 2
+
+    with pytest.raises(ValueError):
+
+        @click.command()
+        @click.argument("method", type=click.EnumChoice(MockEnum, use_value=True))
+        def bad_cli(method):
+            click.echo(method)
+
+    @click.command()
+    @click.argument("method", type=click.EnumChoice(MockEnum, use_value=False))
+    def cli(method):
+        click.echo(method)
+
+    result = runner.invoke(cli, ["foo"])
+    assert not result.exception, "Values do not need to be unique if using enum names"
+    assert result.output == "MockEnum.foo\n"
+
+
 def test_datetime_option_default(runner):
     @click.command()
     @click.option("--start_date", type=click.DateTime())