What's the most idiomatic way to make a collection of enum.Enum in Python?

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)

Answer 1

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.yellow.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()

---

As of Python 3.9 you can add classmethod to property:

    @classmethod
    @property
    def primary_colors(self):
        return self.red, self.blue, self.yellow

which allows calling primary_colors on the enum and not just the member:

>>> Color.primary_colors
(Color.red, Color.blue, Color.yellow)