JohnD
JohnD

Reputation: 486

How do I use typing to force using a specific enum as an argument in python

I’m trying to write a function that would only except a specific enum that each value is a string so the function could not except a string but would get the wanted value from the enum

from enum import Enum

class Options(Enum):
    A = "a"
    B = "b"
    C = "c"

def some_func(option: Options):
    # some code
    return option

The problem I’m experiencing is that if I check the type I get this instead of a string:

>>> type(Options.A)
<enum 'Options'>

I would like to have it return this:

>>> type(Options.A)
<class 'str'>

Any idea how I can implement this so it would work the way I intended? Thanks in advance 😁

Upvotes: 0

Views: 2488

Answers (2)

Ethan Furman
Ethan Furman

Reputation: 69041

>>> type(Options.A)

is always going to return

<enum 'Options'>

because Options.A is an <enum 'Options'> member. However, if you want

>>> isinstance(Options.A, str)

to be

True

then you need to mix in the str type:

class Options(str, Enum):
    A = "a"
    B = "b"
    C = "c"

NB If you mix in the str type, then your Enum members become directly comparable with str:

>>> Options.A == 'a'
True

Upvotes: 1

MEDZ
MEDZ

Reputation: 2295

You can cast the value of the enum attribute:

from enum import Enum

class Options(Enum):
    A = "a"
    B = "b"
    C = "c"


str_enum = str(Options.A)
print(type(str_enum))

Upvotes: 0

Related Questions