Efficent way to find index of member in string enum-CodePudding

I have a constant Enum class that looks something like this:

class Animals(Enum):
    Dog= 'dog'
    Cat= 'cat'
    Chicken = 'chicken'
    Horse = 'horse'

I need to find a simple and efficient way to find the index of one of the members of the Enum. so I came up with the following oneliner:

list(Animals).index(Animals.Chicken)

output:

The problem is, parsing to a list and searching it again is not efficient enough, and I can't change the constant. It feels like there should be a simple solution that I'm missing.

CodePudding user response：

You have a few options:

Use the index values as the values of your Enum, or perhaps IntEnum

class Animals(Enum):
    Dog= 0
    Cat= 1
    Chicken = 2
    Horse = 3
Animals.Chicken.value #returns 2

If you want to keep the textual description too, then subclass Enum, and add the needed attributes (see the Planet example in the docs):

class Animals(Enum):
    Dog = ('dog',0)
    Cat = ('cat',1)
    Chicken = ('chicken',2)
    Horse = ('horse',3)
    def __init__(self, textval, idx):
        self.textval = textval
        self.idx = idx
Animals.Chicken.value #returns ('chicken', 2)
Animals.Chicken.textval #returns 'chicken'
Animals.Chicken.idx #returns 2

CodePudding user response：

Since Enum member names do not start with numbers, you can simply make your Enum subclass' _member_map_ attribute, a dict that stores the member name-to-value mapping, to also store the index number-to-value mapping.

This can be done by subclassing Enum's metaclass, EnumMeta, and adding indices as new keys to all corresponding values for the _member_map_ attribute, in the class object returned by the original __new__ method:

from enum import Enum, EnumMeta

class IndexableEnumType(EnumMeta):
    def __new__(metacls, cls, bases, classdict, **kwargs):
        enum_class = super().__new__(metacls, cls, bases, classdict, **kwargs)
        for i, v in enumerate(list(enum_class._member_map_.values())):
            enum_class._member_map_[i] = v
        return enum_class

So that:

class Animals(Enum, metaclass=IndexableEnumType):
    Dog= 'dog'
    Cat= 'cat'
    Chicken = 'chicken'
    Horse = 'horse'

print(Animals[2])
print(Animals['Chicken'])
print(Animals.Chicken)

would output:

Animals.Chicken
Animals.Chicken
Animals.Chicken

Since dict lookups costs O(1) in average time complexity, this is much more efficient than your current solution of calling list.index, which costs O(n) in time complexity.

Demo: https://replit.com/@blhsing/AdventurousVapidCollaborativesoftware