I have an enum for colors. I wish to add a helper method "toRGB()" to the enum class that converts an instance of the enum to an RGB object. As an optimization, I wished to create the dictionary once as a static variable. However, the correct syntax seems to elude me.
Can anyone suggest the right way to do this?
from enum import Enum
class RGB:
def __init__(self, r, g, b):
pass
class Color(Enum):
RED = 0
GREEN = 1
__tbl = {
RED: RGB(1, 0, 0),
GREEN: RGB(0, 1, 0)
}
def toRGB(self):
return self.__class__.__tbl[self.value]
c = Color.RED
print(c.toRGB())
I get the following error:
Traceback (most recent call last):
File "C:/Users/user/Desktop/test.py", line 20, in <module>
print(c.toRGB())
File "C:/Users/user/Desktop/test.py", line 17, in toRGB
return self.__class__.__tbl[self.value]
TypeError: 'Color' object does not support indexing
Non-method attributes become enum members (even tbl
). You can use a keyword argument instead:
class Color(Enum):
RED = 0
GREEN = 1
def toRGB(self, tbl={
RED: RGB(1, 0, 0),
GREEN: RGB(0, 1, 0)
}):
return tbl[self.value]
Alternatively, you can define the attribute after class creation:
class Color(Enum):
RED = 0
GREEN = 1
def toRGB(self):
return self._tbl[self]
Color._tbl = {
Color.RED: RGB(1, 0, 0),
Color.GREEN: RGB(0, 1, 0)
}