Linking python enums

Question

I have some Enum classes that I want to link up recursively. Looking at the line color = ItemColors[Items(value)._name_].value.value, it seems a bit clunky. Am I misunderstanding something about Enum usage? Is there a better way to do it?

class Colors(Enum):
    white = (255, 255, 255)
    black = (0, 0, 0)
    red = (255, 0, 0)
    green = (0, 255, 0)
    blue = (0, 0, 255)
    yellow = (255, 255, 0)


class ItemColors(Enum):
    empty = Colors.white
    food = Colors.green
    poison = Colors.red
    agent = Colors.blue


class Items(Enum):
    empty = 0
    food = 1
    poison = 2
    agent = 3

items = [0, 0, 3, 2]


def get_item_color(item):
    color = ItemColors[Items(item)._name_].value.value
    return color

Asked By: tsorn

||

Source

Answer 1

Enums are for defining enumerated values. For associating values (items to colors) you can use a map:

class Colors(Enum):
    white = (255, 255, 255)
    black = (0, 0, 0)
    red = (255, 0, 0)
    green = (0, 255, 0)
    blue = (0, 0, 255)
    yellow = (255, 255, 0)

items = [0, 0, 3, 2]

item_colors = {0: Colors.white,
               1: Colors.green,
               2: Colors.red,
               3: Colors.blue}

def get_item_color(item):
    return item_colors[item]

Answered By: brni

Answer 2

Building on @brni’s answer, I came up with a smoother solution:

class Colors(Enum):
    white = (255, 255, 255)
    black = (0, 0, 0)
    red = (255, 0, 0)
    green = (0, 255, 0)
    blue = (0, 0, 255)
    yellow = (255, 255, 0)


class Items(Enum):
    empty = 0
    food = 1
    poison = 2
    agent = 3

item_colors = {
    Items.empty: Colors.white,
    Items.food: Colors.green,
    Items.poison: Colors.red,
    Items.agent: Colors.blue
}

items = [0, 0, 3, 2]


def get_item_color(item):
    color = item_colors[Items(item)].value
    return color

for item in items:
    print(get_item_color(item))

Answered By: tsorn

Answer 3

A better solution is to redefine __new__ with the behavior you desire:

from enum import Enum


class ColorEnum(Enum):
    def __new__(cls, value: int, color: tuple):
        obj = object.__new__(cls)
        obj._value_ = value
        obj.color = color
        return obj


class Items(ColorEnum):
    EMPTY = 0, (255, 255, 255)
    FOOD = 1, (0, 255, 0)
    POSION = 2, (255, 0, 0)
    AGENT = 3, (0, 0, 255)

In the above implementation an extra attribute color is added to each member, however they do not affect reverse lookup (only value is used). See the below example:

>>> print(Items.FOOD.value)
1

>>> print(Items.FOOD.color)
(0, 255, 0)

>>> print(Items(1))
Items.FOOD

>>> print(Items((0, 255, 0)))
ValueError: (0, 255, 0) is not a valid Items

Note: if reverse lookup is desired for both values I recommend MultiValueEnum from the aenum module.

Answered By: Teejay Bruno

Linking python enums

Question:

Answers: