1

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)
1

1 Answer 1

9

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)
5
  • Ahhh, that's exactly the sort of thing I needed, thanks! Commented Sep 14, 2017 at 20:21
  • 1
    @WayneWerner: After thinking about it, I renamed the primary property to is_primary. Commented Sep 15, 2017 at 6:39
  • Color.primary_colors returns a property object, not the actual tuple. It only works on specific enum values, like Color.red.primary_colors Commented May 22, 2023 at 13:00
  • @IvanVirabyan: Thanks -- I should have tested both code paths. Answer updated for 3.9+ where it can work as I thought it should. :-) Commented May 22, 2023 at 16:01
  • Thanks, yeah it works that way, although mypy complains about having both classmethod and property. I went with just classmethod. Commented Jun 15, 2023 at 13:16

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.