Python dictionary with enum as key

Question

Let's say I have an enum

class Color(Enum):
    RED = "RED"
    GREEN = "GREEN"
    BLUE = "BLUE"

I wanted to create a ColorDict class that works as a native python dictionary but only takes the Color enum or its corresponding string value as key.

d = ColorDict() # I want to implement a ColorDict class such that ...

d[Color.RED] = 123
d["RED"] = 456  # I want this to override the previous value
d[Color.RED]    # ==> 456
d["foo"] = 789  # I want this to produce an KeyError exception

What's the "pythonic way" of implementing this ColorDict class? Shall I use inheritance (overriding python's native dict) or composition (keep a dict as a member)?

Answer 1

from enum import Enum

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

# Using enum values as key (with example of type hinting)
color_dict: dict[Color, str] = {
    Color.RED: "red",
    Color.GREEN: "green",
    Color.BLUE: "blue"
}

# Accessing values using enum keys
print(color_dict[Color.RED])  # Outputs: red