Python enum iteration over subset

user_6658749 0
Viewed 3881

I would like to iterate over a subset of the following enum

class Items(enum.Enum):
    item1 = 0
    item2 = 1
    item3 = 2
    item4 = 3
    item5 = 4
    itm66 = 5
    item7 = 6
    item8 = 7

Say I want to:

for item in (Items.item1, Items.item2, Items.item3, Items.item4):
        print(item.value)

is there a shortcut? or do I need to list each item to iterate over?

enums python
4 Answers

Using itertools.islice you can iterate through a slice of your Enum class

from enum import Enum
from itertools import islice

class Items(Enum):
    item1 = 0
    item2 = 1
    item3 = 2
    item4 = 3
    item5 = 4
    itm66 = 5
    item7 = 6
    item8 = 7

for i in islice(Items, 4):
    print(i.value)

# 0 
# 1 
# 2
# 3
user_102556520

Python enums can have methods. I'd suggest you write a method that returns an iterable. Probably a set, in this case:

class Items(enum.Enum):
    item1 = 0
    item2 = 1
    item3 = 2
    item4 = 3
    item5 = 4
    itm66 = 5
    item7 = 6
    item8 = 7

    @classmethod
    def the_best_ones(cls):
        return cls.item1, cls.item2, cls.item3, cls.item4

Then:

for item in Items.the_best_ones():
    print(item.value)
user_40290140

There is nothing built-in to Enum to iterate over a subset, but since you can use list on them:

>>> list(Items)[:4]
[<Items.item1: 0>, <Items.item2: 1>, <Items.item3: 2>, <Items.item4: 3>]
user_2088800

do I need to list each item to iterate over?

No - Enum classes are iterables, each value has attributes .name and .value, so you can filter on either as you see fit. For example:

for v in Items:
    if v.value > 3:
       break
    print(v.name, v.value)

=> 
item1 0
item2 1
item3 2
item4 3

is there a shortcut?

Depends on what you want to do. Here are a few ways you could use the fact that Enums are iterables:

  • sorted(Items, key=lambda v: v.value)[:4] => get the first 4 elements as sorted by their ordinal value
  • filter(lambda e: e.value < 5, Items) => get all items with an ordinal value < 5
  • {v.name: v.value for v in Items if v.value < 4} => get a dictionary of all name/value pairs given a filter

etc.

Note

According to the documentation you can get the Enums ordered dictionary by Items.__members__ and thus you could use .items on that to get the key/value mappings. However what you get as the value (in the dictionary) is in fact an object instance that has the .name and .value attributes.

Items.__members__
=> 
mappingproxy({'item1': <Items.item1: 0>,
              'item2': <Items.item2: 1>,
              'item3': <Items.item3: 2>,
              'item4': <Items.item4: 3>)
# so you could write
for k, v in Items.__members__.items():
    if v > 3:
       break
    print(k, v.value)

However I find the first method more intuitive.

user_8902420
Related