How to add member subsets to a Python Enum?

Question

Say I have a Python Enum like this:

from enum import Enum

class FRUIT(Enum):
    APPLE = 1
    BANANA = 2
    LEMON = 3
    ORANGE = 4

and I want to "subset" these in useful ways, like be able to say:

if fruit in FRUIT.CITRUS:
    make_juice()

where somewhere I have defined: CITRUS = {LEMON, ORANGE}.

I'd like to keep the subsets as attributes of the main Enum because it keeps use of the subset in context.

I know I can do something like this but I'd strongly prefer to avoid the method-call notation. Also it seems wasteful to reconstruct the set each time its needed:

@classmethod
def CITRUS(cls):
    return {cls.LEMON, cls.ORANGE}

Is there a way to add a class attribute after the Enum meta-class has done its work without breaking things?

Answer 1

Since CITRUS is not meant to be a fruit itself, but a type of fruit, it makes more sense to create a separate Enum subclass with the types of fruit as the members:

class FRUIT_TYPE(Enum):
    CITRUS = {FRUIT.LEMON, FRUIT.ORANGE}

so that you can do something like:

fruit = FRUIT.LEMON
if fruit in FRUIT_TYPE.CITRUS.value:
     make_juice()

Having to use FRUIT_TYPE.CITRUS.value to check for membership looks cumbersome, however. To allow membership check of FRUIT_TYPE.CITRUS itself you can make FRUIT_TYPE a subclass of set as well:

class FRUIT_TYPE(set, Enum):
    CITRUS = {FRUIT.LEMON, FRUIT.ORANGE}

so that you can do the following instead:

fruit = FRUIT.LEMON
if fruit in FRUIT_TYPE.CITRUS:
     make_juice()

Answer 2

If you don't need specific values for the Enum values, it is likely you can get all the functionality you need by using enum.IntFlag enumerations rather than simply Enum.

Just declare your Enum class as an IntFlag, and you are free to use &, | and possibly other bitwise operators to have the behavior you need:

In [1]: import enum                                                                                                                  

In [2]: class Fruit(enum.IntFlag): 
   ...:     APPLE = 1 
   ...:     BANANA = 2 
   ...:     LEMON = 4 
   ...:     ORANGE = 8 
In [4]: CITRUS = Fruit.LEMON | Fruit.ORANGE                                                                                          


In [6]: for fruit in Fruit: 
   ...:     if fruit & CITRUS: 
   ...:         print(f"making_juice({fruit.name})") 
   ...:          
   ...:   

This does not allow for interaction directly on "CITRUS", and requires a filter pattern, like I used above.

However, as recently as a few weeks ago, I wanted exactly this feature, and could implement it as an __iter__ method doing this filter straight in the Enum class:

    def __iter__(self):
        for element in self.__class__:
            if self & element:
                yield element

If we simply plug that in the above enumeration:

In [8]: class Fruit(enum.IntFlag): 
   ...:     APPLE = 1 
   ...:     BANANA = 2 
   ...:     LEMON = 4 
   ...:     ORANGE = 8 
   ...:      
   ...:     def __iter__(self): 
   ...:         for element in self.__class__: 
   ...:             if self & element: 
   ...:                 yield element 
   ...:                                                                                                                              

In [9]: CITRUS = Fruit.LEMON | Fruit.ORANGE                                                                                          

In [10]: for fruit in CITRUS: 
    ...:     print (fruit.name) 
    ...:                                                                                                                             
LEMON
ORANGE

The __iter__ does not conflict with iterating on the Fruit class itself, as that uses the __iter__ method in the EnumMeta metaclass and, as can be seen, will be correctly called by "ORed" subsets of the enums. Which means that if you need them you just have to write proper __len__ and __contains__ methods to have all the features you'd expect from subsets.

I am using this code in a personal project, and it works like a charm:

https://github.com/jsbueno/terminedia/blob/9714d6890b8336678cd10e0c6275f56392e409ed/terminedia/values.py#L51

(Although, right now the "unicode_effects" declared just below the enumeration there is an ordinary set, now that you mentioned it I think I will just write __contains__ and use that instead of the set.)

Answer 3

UPDATE: One other thing to consider, which I think I like better, on balance, is to define subset membership as an attribute of each Enum member, like:

fruit = FRUIT.ORANGE  # ---or whatever, probably in far away code---
...
if fruit.is_citrus:
    make_juice()

These can be defined as @propertys on the class and don't suffer from the mutability problems mentioned below.

class FRUIT(Enum):
    APPLE = 1
    BANANA = 2
    LEMON = 3
    ORANGE = 4

    @property
    def is_citrus(self):
        return self in frozenset((FRUIT.LEMON, FRUIT.ORANGE))

---

---

Thanks to the other respondents who all contributed very useful points of view. Here's what I ended up doing after considering the other answers, followed by my rationale:

from enum import Enum

class FRUIT(Enum):
    APPLE = 1
    BANANA = 2
    LEMON = 3
    ORANGE = 4

FRUIT.CITRUS_TYPES = frozenset((FRUIT.LEMON, FRUIT.ORANGE))

This works fine and (suprisingly to me) doesn't break any of the other Enum behaviors:

# ---CITRUS_TYPES subset has desired behavior---
>>> FRUIT.LEMON in FRUIT.CITRUS_TYPES
True
>>> FRUIT.APPLE in FRUIT.CITRUS_TYPES
False
>>> "foobar" in FRUIT.CITRUS_TYPES
False

# ---CITRUS_TYPES has not become a member of FRUIT enum---
>>> tuple(FRUIT)
(FRUIT.APPLE: 1>, <FRUIT.BANANA: 2>, <FRUIT.LEMON: 3>, <FRUIT.ORANGE: 4>)
>>> FRUIT.APPLE in FRUIT
True
>>> FRUIT.CITRUS_TYPES in FRUIT
DeprecationWarning: using non-Enums in containment checks will raise TypeError in Python 3.8
False

# ---CITRUS_TYPES not reported by dir(FRUIT)---
>>> dir(FRUIT)
['APPLE', 'BANANA', 'LEMON', 'ORANGE', '__class__', '__doc__', '__members__', '__module__']

# ---But it does appear on FRUIT.__dict__---
FRUIT.__dict__ == {
    '_generate_next_value_': <function Enum._generate_next_value_ at 0x1010e9268>, 
    '__module__': '__main__',
    '__doc__': 'An enumeration.',
    '_member_names_': ['APPLE', 'BANANA', 'LEMON', 'ORANGE'],
    '_member_map_': OrderedDict([
        ('APPLE', <FRUIT.APPLE: 1>),
        ('BANANA', <FRUIT.BANANA: 2>),
        ('LEMON', <FRUIT.LEMON: 3>),
        ('ORANGE', <FRUIT.ORANGE: 4>)
    ]),
    '_member_type_': <class 'object'>,
    '_value2member_map_': {
        1: <FRUIT.APPLE: 1>,
        2: <FRUIT.BANANA: 2>,
        3: <FRUIT.LEMON: 3>,
        4: <FRUIT.ORANGE: 4>,
    },
    'APPLE': <FRUIT.APPLE: 1>,
    'BANANA': <FRUIT.BANANA: 2>,
    'LEMON': <FRUIT.LEMON: 3>,
    'ORANGE': <FRUIT.ORANGE: 4>,
    '__new__': <function Enum.__new__ at 0x1010e91e0>,
    'CITRUS_TYPES': frozenset({<FRUIT.LEMON: 3>, <FRUIT.ORANGE: 4>})
}

So it appears to store CITRUS_TYPES on the class, but hides it from dir() for whatever reason.

It does have a vulnerability though, in that the added attribute is mutable, like any other class attribute; if some part of the client code assigns to FRUIT.CITRUS_TYPES, FRUIT would not complain and that would of course break things. This behavior is unlike that of an Enum member, which raises AttributeError on attempted assignment.

I thought this might be remedied by making it a classproperty, which I ended up trying, but my early attempts did not prevent mutation. A more sophisticated classproperty implementation described there might work, but I ended up be satisfied with the simple approach above for now.

@blhsing raises an interesting question on whether such an attribute on FRUIT makes sense. I see his point, and may yet adopt his view, but my current view is that localizing "fruit-related" characteristics to a single imported name is best for me. One could consider FRUIT a rigorous set of fruit-types and subsets of fruit to therefore be a distinct set. I find that rigor unrewarding for my current purposes and prefer to think of FRUIT as more of a collection of related constant values, including both members and subsets. YMMV of course. Like I say, I may yet adopt his view.