From 828277b1e57b39b204171b256fef2e7995c468de Mon Sep 17 00:00:00 2001 From: Sagi Shadur Date: Thu, 10 Mar 2022 13:20:27 +0200 Subject: [PATCH 1/3] Add `EnumChoice` parameter type --- src/click/__init__.py | 1 + src/click/types.py | 18 ++++++++++ tests/test_basic.py | 52 +++++++++++++++++++++++++++ tests/test_info_dict.py | 22 ++++++++++++ tests/test_normalization.py | 21 +++++++++++ tests/test_options.py | 71 +++++++++++++++++++++++++++++++++++++ 6 files changed, 185 insertions(+) diff --git a/src/click/__init__.py b/src/click/__init__.py index 33080c0a6..bb5f551f9 100644 --- a/src/click/__init__.py +++ b/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 diff --git a/src/click/types.py b/src/click/types.py index aab0656ed..06ac36050 100644 --- a/src/click/types.py +++ b/src/click/types.py @@ -1,3 +1,4 @@ +import enum import os import stat import typing as t @@ -322,6 +323,23 @@ def shell_complete( return [CompletionItem(c) for c in matched] +class EnumChoice(Choice): + def __init__(self, enum_type: t.Type[enum.Enum], case_sensitive: bool = True): + super().__init__( + choices=[element.name for element in enum_type], + case_sensitive=case_sensitive, + ) + self.enum_type = enum_type + + def convert( + self, value: t.Any, param: t.Optional["Parameter"], ctx: t.Optional["Context"] + ) -> t.Any: + value = super().convert(value=value, param=param, ctx=ctx) + if value is None: + return None + return self.enum_type[value] + + class DateTime(ParamType): """The DateTime type converts date strings into `datetime` objects. diff --git a/tests/test_basic.py b/tests/test_basic.py index d68b96299..517fcc324 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,3 +1,4 @@ +import enum import os from itertools import chain @@ -6,6 +7,15 @@ import click +class MyEnum(enum.Enum): + """Dummy enum for unit tests.""" + + ONE = "one" + TWO = "two" + THREE = "three" + ONE_ALIAS = ONE + + def test_basic_functionality(runner): @click.command() def cli(): @@ -403,6 +413,48 @@ def cli(method): assert "{foo|bar|baz}" in result.output +def test_enum_choice_option(runner): + @click.command() + @click.option("--number", type=click.EnumChoice(MyEnum)) + def cli(number): + click.echo(number) + + result = runner.invoke(cli, ["--number=ONE"]) + assert not result.exception + assert result.output == "MyEnum.ONE\n" + + result = runner.invoke(cli, ["--number=meh"]) + assert result.exit_code == 2 + assert ( + "Invalid value for '--number': 'meh' is not one of 'ONE', 'TWO', 'THREE'." + in result.output + ) + + result = runner.invoke(cli, ["--help"]) + assert "--number [ONE|TWO|THREE]" in result.output + + +def test_enum_choice_argument(runner): + @click.command() + @click.argument("number", type=click.EnumChoice(MyEnum)) + def cli(number): + click.echo(number) + + result = runner.invoke(cli, ["ONE"]) + assert not result.exception + assert result.output == "MyEnum.ONE\n" + + result = runner.invoke(cli, ["meh"]) + assert result.exit_code == 2 + assert ( + "Invalid value for '{ONE|TWO|THREE}': 'meh' is not one of 'ONE', " + "'TWO', 'THREE'." in result.output + ) + + result = runner.invoke(cli, ["--help"]) + assert "{ONE|TWO|THREE}" in result.output + + def test_datetime_option_default(runner): @click.command() @click.option("--start_date", type=click.DateTime()) diff --git a/tests/test_info_dict.py b/tests/test_info_dict.py index b58ad6efb..a478be810 100644 --- a/tests/test_info_dict.py +++ b/tests/test_info_dict.py @@ -1,8 +1,11 @@ +import enum + import pytest import click.types # Common (obj, expect) pairs used to construct multiple tests. + STRING_PARAM_TYPE = (click.STRING, {"param_type": "String", "name": "text"}) INT_PARAM_TYPE = (click.INT, {"param_type": "Int", "name": "integer"}) BOOL_PARAM_TYPE = (click.BOOL, {"param_type": "Bool", "name": "boolean"}) @@ -91,6 +94,15 @@ ) +class MyEnum(enum.Enum): + """Dummy enum for unit tests.""" + + ONE = "one" + TWO = "two" + THREE = "three" + ONE_ALIAS = ONE + + @pytest.mark.parametrize( ("obj", "expect"), [ @@ -115,6 +127,16 @@ }, id="Choice ParamType", ), + pytest.param( + click.EnumChoice(MyEnum), + { + "param_type": "EnumChoice", + "name": "choice", + "choices": ["ONE", "TWO", "THREE"], + "case_sensitive": True, + }, + id="Choice ParamType", + ), pytest.param( click.DateTime(["%Y-%m-%d"]), {"param_type": "DateTime", "name": "datetime", "formats": ["%Y-%m-%d"]}, diff --git a/tests/test_normalization.py b/tests/test_normalization.py index 32df098e9..4d982a548 100644 --- a/tests/test_normalization.py +++ b/tests/test_normalization.py @@ -1,9 +1,20 @@ +import enum + import click CONTEXT_SETTINGS = dict(token_normalize_func=lambda x: x.lower()) +class MyEnum(enum.Enum): + """Dummy enum for unit tests.""" + + ONE = "one" + TWO = "two" + THREE = "three" + ONE_ALIAS = ONE + + def test_option_normalization(runner): @click.command(context_settings=CONTEXT_SETTINGS) @click.option("--foo") @@ -26,6 +37,16 @@ def cli(choice): assert result.output == "Foo\n" +def test_enum_choice_normalization(runner): + @click.command(context_settings=CONTEXT_SETTINGS) + @click.option("--choice", type=click.EnumChoice(MyEnum)) + def cli(choice): + click.echo(choice) + + result = runner.invoke(cli, ["--CHOICE", "ONE"]) + assert result.output == "MyEnum.ONE\n" + + def test_command_normalization(runner): @click.group(context_settings=CONTEXT_SETTINGS) def cli(): diff --git a/tests/test_options.py b/tests/test_options.py index 3beff11ba..10ef16e89 100644 --- a/tests/test_options.py +++ b/tests/test_options.py @@ -1,3 +1,4 @@ +import enum import os import re @@ -7,6 +8,15 @@ from click import Option +class MyEnum(enum.Enum): + """Dummy enum for unit tests.""" + + ONE = "one" + TWO = "two" + THREE = "three" + ONE_ALIAS = ONE + + def test_prefixes(runner): @click.command() @click.option("++foo", is_flag=True, help="das foo") @@ -570,6 +580,67 @@ def cmd(foo): assert result.output == "Apple\n" +def test_missing_enum_choice(runner): + @click.command() + @click.option("--foo", type=click.EnumChoice(MyEnum), required=True) + def cmd(foo): + click.echo(foo) + + result = runner.invoke(cmd) + assert result.exit_code == 2 + error, separator, choices = result.output.partition("Choose from") + assert "Error: Missing option '--foo'. " in error + assert "Choose from" in separator + assert "ONE" in choices + assert "TWO" in choices + assert "THREE" in choices + assert "ONE_ALIAS" not in choices + + +def test_case_insensitive_enum_choice(runner): + @click.command() + @click.option("--foo", type=click.EnumChoice(MyEnum, case_sensitive=False)) + def cmd(foo): + click.echo(foo) + + result = runner.invoke(cmd, ["--foo", "one"]) + assert result.exit_code == 0 + assert result.output == "MyEnum.ONE\n" + + result = runner.invoke(cmd, ["--foo", "tHREE"]) + assert result.exit_code == 0 + assert result.output == "MyEnum.THREE\n" + + result = runner.invoke(cmd, ["--foo", "Two"]) + assert result.exit_code == 0 + assert result.output == "MyEnum.TWO\n" + + @click.command() + @click.option("--foo", type=click.EnumChoice(MyEnum)) + def cmd2(foo): + click.echo(foo) + + result = runner.invoke(cmd2, ["--foo", "one"]) + assert result.exit_code == 2 + + result = runner.invoke(cmd2, ["--foo", "tHREE"]) + assert result.exit_code == 2 + + result = runner.invoke(cmd2, ["--foo", "TWO"]) + assert result.exit_code == 0 + + +def test_case_insensitive_enum_choice_returned_exactly(runner): + @click.command() + @click.option("--foo", type=click.EnumChoice(MyEnum, case_sensitive=False)) + def cmd(foo): + click.echo(foo) + + result = runner.invoke(cmd, ["--foo", "ONE"]) + assert result.exit_code == 0 + assert result.output == "MyEnum.ONE\n" + + def test_option_help_preserve_paragraphs(runner): @click.command() @click.option( From ac6e2a1877fb307579fa6ffa208a87aedb8d5cb3 Mon Sep 17 00:00:00 2001 From: Sagi Shadur Date: Thu, 10 Mar 2022 13:56:54 +0200 Subject: [PATCH 2/3] Fix tests id in test_info_dict.py --- tests/test_info_dict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_info_dict.py b/tests/test_info_dict.py index a478be810..f6b1ca35f 100644 --- a/tests/test_info_dict.py +++ b/tests/test_info_dict.py @@ -135,7 +135,7 @@ class MyEnum(enum.Enum): "choices": ["ONE", "TWO", "THREE"], "case_sensitive": True, }, - id="Choice ParamType", + id="EnumChoice ParamType", ), pytest.param( click.DateTime(["%Y-%m-%d"]), From 28e019089343d860e2e2ab6989894cf1b40327db Mon Sep 17 00:00:00 2001 From: Sagi Buchbinder Shadur Date: Mon, 21 Aug 2023 08:55:20 -0400 Subject: [PATCH 3/3] Fix pre-commit issues --- src/click/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/click/types.py b/src/click/types.py index aca569833..379e67cc8 100644 --- a/src/click/types.py +++ b/src/click/types.py @@ -345,7 +345,7 @@ def __init__(self, enum_type: t.Type[enum.Enum], case_sensitive: bool = True): self.enum_type = enum_type def convert( - self, value: t.Any, param: t.Optional["Parameter"], ctx: t.Optional["Context"] + self, value: t.Any, param: t.Optional[Parameter], ctx: t.Optional[Context] ) -> t.Any: value = super().convert(value=value, param=param, ctx=ctx) if value is None: