Python Enum support for click.Choice

[devraj]

I can work around it as described here it would be great if Choice supported Python Enum.

Unless of course I have missed something completely fundamental :-)

[mitsuhiko]

What exactly is implied with supporting an enum?

[devraj]

To be able to use Enum instead of tuples for Choice values, i.e define it as such

from enum import Enum, unique

@unique
class ConfigFormat(Enum):
    yaml = 0
    json = 1
    plist = 2

and then use it in the decorator as follows

from . import const

@dispatch.command()
@click.option('--output', '-o', type=click.Choice(const.ConfigFormat),
                help='Sets default output format for configuration files')
@pass_project
def init(project, output, force):
    """Initialises a managed schema"""
    click.echo(project.schema_home)

[aldanor]

enum is only available in Python 3 though (in Python 2 you have to use an external backport).

[TAGC]

I was being facetious, sorry. It would be a neat feature but it's nothing essential by any means. Loving this tool by the way.

[allanlewis]

I'm using something like this in my code, where I'm using Python 2.7 with enum34==1.1.6:

@click.option(
    '--enum-val', type=click.Choice(MyEnum.__members__),
    callback=lambda c, p, v: getattr(MyEnum, v) if v else None)

The callback provides enum_val as the actual enumeration instance rather than a string, or None if the option wasn't given.

Perhaps this could be wrapped into another decorator, like click.enum_option:

@click.enum_option('--enum-val', enum=MyEnum)

UPDATE: In review, others thought that callback was rather ugly, which it is, so I've removed it. If it was inside click, it might be OK.

[skycaptain]

The easiest way I can think of, is to use a custom type (Although I think this would be a nice feature for click):

class EnumType(click.Choice):
    def __init__(self, enum):
        self.__enum = enum
        super().__init__(enum.__members__)

    def convert(self, value, param, ctx):
        return self.__enum[super().convert(value, param, ctx)]

You might overwrite get_metavar to compute the metavar from class name (since the complete choice list might be to long for a cli printout):

...
    def get_metavar(self, param):
        # Gets metavar automatically from enum name
        word = self.__enum.__name__

        # Stolen from jpvanhal/inflection
        word = re.sub(r"([A-Z]+)([A-Z][a-z])", r'\1_\2', word)
        word = re.sub(r"([a-z\d])([A-Z])", r'\1_\2', word)
        word = word.replace("-", "_").lower().split("_")

        if word[-1] == "enum":
            word.pop()

        return ("_".join(word)).upper()
...

Since Enums are sometimes written uppercase, feel free to write a case insensitive version of the above code. E.g. one could start off with

class EnumType(click.Choice):
    def __init__(self, enum, casesensitive=True):
        if isinstance(enum, tuple):
            choices = (_.name for _ in enum)
        elif isinstance(enum, EnumMeta):
            choices = enum.__members__
        else:
            raise TypeError("`enum` must be `tuple` or `Enum`")

        if not casesensitive:
            choices = (_.lower() for _ in choices)

        self.__enum = enum
        self.__casesensitive = casesensitive

        # TODO choices do not have the save order as enum
        super().__init__(list(sorted(set(choices))))

    def convert(self, value, param, ctx):
        if not self.__casesensitive:
            value = value.lower()

        value = super().convert(value, param, ctx)

        if not self.__casesensitive:
            return next(_ for _ in self._EnumType__enum if _.name.lower() ==
                            value.lower())
        else:
            return next(_ for _ in self._EnumType__enum if _.name == value)

    def get_metavar(self, param):
        word = self.__enum.__name__

        # Stolen from jpvanhal/inflection
        word = re.sub(r"([A-Z]+)([A-Z][a-z])", r'\1_\2', word)
        word = re.sub(r"([a-z\d])([A-Z])", r'\1_\2', word)
        word = word.replace("-", "_").lower().split("_")

        if word[-1] == "enum":
            word.pop()

        return ("_".join(word)).upper()

This way you can either add a complete Enum or you can just use some values as a choice (listed as a tuple).

[aomader]

A bit old but I thought I add the nice idea of using str Enums for options (which is also great for old-school str-based settings w.r.t. comparisons).

class ChoiceType(click.Choice):
    def __init__(self, enum):
        super().__init__(map(str, enum))
        self.enum = enum

    def convert(self, value, param, ctx):
        value = super().convert(value, param, ctx)
        return next(v for v in self.enum if str(v) == value)

class Choice(str, Enum):
    def __str__(self):
        return str(self.value)

class MyChoice(Choice):
    OPT_A = 'opt-a'
    OPT_B = 'opt-b'

@click.option('--choice', type=MyChoiceType(MyChoice),
              default=MyChoice.OPT_B)
def func(choice):
    assert choice in ('opt-a', 'opt-b')
    assert choice in (MyChoice.OPT_A, MyChoice.OPT_B)
    

[gazpachoking]

Here's my take on supporting this:

class EnumChoice(click.Choice):
    def __init__(self, enum, case_sensitive=False, use_value=False):
        self.enum = enum
        self.use_value = use_value
        choices = [str(e.value) if use_value else e.name for e in self.enum]
        super().__init__(choices, case_sensitive)

    def convert(self, value, param, ctx):
        if value in self.enum:
            return value
        result = super().convert(value, param, ctx)
        # Find the original case in the enum
        if not self.case_sensitive and result not in self.choices:
            result = next(c for c in self.choices if result.lower() == c.lower())
        if self.use_value:
            return next(e for e in self.enum if str(e.value) == result)
        return self.enum[result]

Allows using either the names or values of the enum items based on the use_value parameter. Should work whether values are strings or not.

[bartekpacia]

I don't know if somebody posted it here, but that's how I dealt with it:

Language = Enum("Language", "pl en")

@click.option("--language", type=click.Choice(list(map(lambda x: x.name, Language)), case_sensitive=False))
def main():
...

[sscherfke]

The recipe above no longer works properly with click 8.

When the help text is generated, an option type's convert() functions is eventually being called converting the string denoting the Enum’s attribute in to the attribute value itself and --my-opt [a|b|c] [default: a] becomes --my-opt [a|b|c] [default: MyEnum.a].

I looked through the code and found no easy way (for library authors) to fix this.

[ahmed-shariff]

This is an ugly hack, for now I am using a slightly modified version of @allanlewis:

class MyEnum(Enum):
    a = "a"
    b = "b"

@click.option("-m", "--method", type=click.Choice(MyEnum.__members__), 
              callback=lambda c, p, v: getattr(MyEnum, v) if v else None, default="a")

[yoyonel]

Proposition

Less ugly solution (but not perfect, seems have some problems with typing (especially inside pycharm)):

from enum import Enum

import click

MyEnum = Enum("my_enum", ("a", "b"))


@click.option(
    "-m", "--method", 
    type=click.Choice(list(map(lambda x: x.name, MyEnum)), case_sensitive=False),
    default="a"
)

Remove the (ugly) usage of __members__ and more concise :-)

Documentation

• Enum HOWTO - Functional API

[allanlewis]

@yoyonel you could replace the lambda with attrgetter('name') 🙂

[yashrathi-git]

This is easy to implement, using custom type

import click
from enum import Enum

class Test(Enum):
    test = "test"
    another_option = "another_option"

class EnumType(click.Choice):
    def __init__(self, enum, case_sensitive=False):
        self.__enum = enum
        super().__init__(choices=[item.value for item in enum], case_sensitive=case_sensitive)

    def convert(self, value, param, ctx):
        converted_str = super().convert(value, param, ctx)
        return self.__enum(converted_str)

@click.option("-m", "--method", type=EnumType(Test), default = "test")
@click.command()
def test(method):
    print(type(method))
    print(method)
    print(method.value)

if __name__ == '__main__':
    test()

Test:

$ python test.py --method wrong
Usage: test.py [OPTIONS]
Try 'test.py --help' for help.

Error: Invalid value for '-m' / '--method': invalid choice: wrong. (choose from test, another_option)


$ python test.py --method test 
<enum 'Test'>
Test.test
test

[ShivKJ]

@yashrathi-git a minor change in the code you provided,

from enum import Enum

from click import Choice

class Test(Enum):
    test = "test"
    another_option = "another_option"

    def __str__(self):
        return self.value


class EnumType(Choice):
    def __init__(self, enum: Enum, case_sensitive=False):
        self.__enum = enum
        super().__init__(choices=[item.value for item in enum], case_sensitive=case_sensitive)

    def convert(self, value, param, ctx):
        if value is None or isinstance(value, Enum)::
            return value

        converted_str = super().convert(value, param, ctx)
        return self.__enum(converted_str)

But I feel this should be done under the hood instead of importing EnumType.

In other words, it would be much more desirable if we could write,

@click.option("-m", "--method", type=Test, default = 'test') # or default = Test.test

[jerluc]

I know this thread is pretty old, but I've been using this pattern for str + Enums, which IMO is pretty clean and intuitive:

import click
import enum


class MyEnum(str, enum.Enum):
    A = "a"
    B = "b"
    C = "c"


@click.command("my-cmd")
@click.argument("which_option", type=click.Choice(MyEnum))
def my_cmd(which_option: MyEnum):
    print(which_option)
    print(type(which_option))


if __name__ == "__main__":
    my_cmd()

$ python click_enum.py

Usage: click_enum.py [OPTIONS] [a|b|c]
Try 'click_enum.py --help' for help.

Error: Missing argument '[a|b|c]'.  Choose from:
	a,
	b,
	c.

$ python click_enum.py a

MyEnum.A
<enum 'MyEnum'>

$ python click_enum.py d

Usage: click_enum.py [OPTIONS] [a|b|c]
Try 'click_enum.py --help' for help.

Error: Invalid value for '[a|b|c]': invalid choice: d. (choose from a, b, c)

[rdbisme]

Hello @jerluc, your pattern works functionally, but it fails type checking

[jerluc]

Hello @jerluc, your pattern works functionally, but it fails type checking

Does it @rdbisme? I don't see any issues using latest mypy (0.931) on Python 3.7:

jerluc@ws ~ $ cat click_enum.py
import click
import enum


class MyEnum(str, enum.Enum):
    A = "a"
    B = "b"
    C = "c"


@click.command("my-cmd")
@click.argument("which_option", type=click.Choice(MyEnum))
def my_cmd(which_option: MyEnum):
    print(which_option)
    print(type(which_option))


if __name__ == "__main__":
    my_cmd()
jerluc@ws ~ $ mypy click_enum.py
Success: no issues found in 1 source file
jerluc@ws ~ $ mypy --version
mypy 0.931

[rdbisme]

Well,

$ cat click_test.py
import click
import enum


class MyEnum(str, enum.Enum):
    A = "a"
    B = "b"
    C = "c"


@click.command("my-cmd")
@click.argument("which_option", type=click.Choice(MyEnum))
def my_cmd(which_option: MyEnum):
    print(which_option)
    print(type(which_option))


if __name__ == "__main__":
    my_cmd()

Still Gives:

click_test.py:12: error: Argument 1 to "Choice" has incompatible type "Type[MyEnum]"; expected "Sequence[str]"
Found 1 error in 1 file (checked 1 source file)

$ python --version
Python 3.9.9

$ python -c "import click; print(click.__version__)"
8.0.3

[dzcode]

I'm at the PyCon sprint now and could work on this

[davidism]

@dzcode thanks, but there's already an open PR for this.

[MicaelJarniac]

@dzcode thanks, but there's already an open PR for this.

#2210?