pythonenumspython-3.9

Python Enum class - Use Enum name argument to select Enum value within a function


I'm using Python 3.9. I have an Enum class like this:

class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

I want to use this enum class in a function with an argument which is a valid Enum name:

def some_func(color_name):
    return Color.color_name.value

So, when I call the function:

>>> some_func(RED)

why does it not work?

Do I need to create a method for the Color class? That seems to defeat the purpose of have it as an Enum class.


Solution

  • 2 issues, one you are passing in RED instead of "RED" where the first is undefined. Small mistake of course and isn't the main issue. Main issue is that passing in just the color name (even if you fix the first thing and correctly pass in a string), you are trying to access enum attributes like dictionaries with just using a string as a key. Enums don't quite work that way. You can use getattr() instead to find the Color Enum type that matches what you passed in.

    i.e. getattr(Color, color_name).value

    Here is your full corrected function

    class Color(Enum):
        RED = 1
        GREEN = 2
        BLUE = 3
    
    def some_func(color_name):
        return getattr(Color, color_name).value
    
    print(some_func("RED"))
    

    However a simpler way than all of this is just to access the enum values directly.

    access the enum classes value directly
    print(Color.RED.value)
    

    Note here the 'RED' is NOT a string, it is an attribute of the enum.

    Note that this will return the value of the enum (1 in this example), not the enum type itself.