What's the most idiomatic way to make a collection of enum.Enum in Python?
https://softwareengineering.stackexchange.com/questions/357405
Question
I have an Enum in Python that looks something like this:
import enum
class Color(enum.Enum):
red = 'red'
blue = 'blue'
yellow = 'yellow'
puce = 'puce'
chartreuse = 'chartreuse'
And I want to have something like:
primary_colors = (Color.red, Color.blue, Color.yellow)
But naturally if I try and stick that line in the enum, Color
is not defined. And if I use (red, blue, yellow)
then it becomes just a tuple of strings. And I can't extend Enum, so I can't do something like:
class PrimaryColor(enum.Enum):
red = 'red'
yellow = 'yellow'
green = 'green'
class Color(enum.Enum, PrimaryColor):
puce = 'puce'
chartreuse = 'chartreuse'
So what's an idiomatic sort of way that I can say that I have this collection of primary colors that's a subset of Color
? Am I limited to something like this?
primary_colors = (Color.red, Color.yellow, Color.blue)
Solution
One of the things to keep in mind about enum.Enum
is that any non-descriptor
attribute is converted into an Enum
member -- so property
and functions/methods are not (converted). This means you can do something like:
class Color(enum.Enum):
red = 'red'
blue = 'blue'
yellow = 'yellow'
puce = 'puce'
chartreuse = 'chartreuse'
@property
def primary_colors(self):
return self.red, self.blue, self.yellow
This allows you to do:
>>> Color.primary_colors
(Color.red, Color.blue, Color.yellow)
If you wanted to, you could add an is_primary()
property as well:
@property
def is_primary(self):
return self in self.primary_colors
And then, if you get a Color
member, do:
if some_Color_member.is_primary:
do_something()