CODEWITHSUNDEEP
pythonpython-3.xlistenumsenumeration
Robert
Robert

Reputation: 515

Iterate or index list based on Enum

I have a small enum list of 4 elements (class Type(Enum): One,Two,Three,Four), and I'm wondering if there is a way to associate the enum with a global list (List = [Type(0),Type(1),Type(2),Type(3)].

I know I can access this list using the syntax List[ Type.Two.value ]. But I'm hoping there is a simple way to lose the .value. It's not a typing problem - I don't mind typing long descriptive variable names. But the list and enum are directly related/associated, so it would be preferable if the list respected the enum as an iterator. Is there a simple way to pull this off?

I'm working in the environment of Blender, building an add-on, if that makes any difference. I'm also very new to python altogether - about a week old, so do not omit obvious points if you have them.

I'm very familiar with C/C++, in case there are any similarities. I was a big fan of creating data library lists and associating iterator types in a way that prevented access to the list using anything other than the index typename associated with it (so you couldn't make a mistake like Cars[my_boat_index])

I appreciate any advice!

Upvotes: 2

Views: 4806

Answers (3)

Ethan Furman
Ethan Furman

Reputation: 69041

Turning an Enum into a list is as easy as:

list(Type)
# [<Type.ONE: 1>, <Type.TWO: 2>, <Type.THREE: 3>, <Type.FOUR: 4>]

Accessing a list using Type members as indices is a little more complicated:

  • by default, member values start at 1
  • but container indices start at 0
  • and members don't have the __index__ and __int__ methods

The easy/simple solution is twofold:

  • use IntEnum instead of Enum
  • start the counting at 0

So:

class Type(IntEnum):
    ZERO = 0
    ONE = 1
    TWO = 2
    THREE = 3

IntEnum members can be used anywhere an int is, so they work fine as index values.

The more complicated solution is to use a regular Enum, add in the methods you need, and subtract 1 when returning the value:

class Type(Enum):
    ONE = 1
    TWO = 2
    THREE = 3
    FOUR = 4

    def __index__(self):
        return self.value - 1

    def __int__(self):
        return self.value - 1

Even if you go with adding the __index__ and __int__ methods, you should still just start counting at zero instead of doing the subtract one trick, as that will cause more problems later when you get used to the way Python works.

Upvotes: 3

Grismar
Grismar

Reputation: 31319

It appears you're looking to restrict access to a list by restricting the type of the index to a specific type (e.g. an Enum).

Although I don't think it's a good idea, you can definitely do it:

from enum import Enum


class MyEnum(Enum):
    ZERO = 0
    ONE = 1
    TWO = 2


class MyList(list):
    def __getitem__(self, item):
        assert isinstance(item, MyEnum)
        return super().__getitem__(item.value)


ml = MyList([1, 2, 3])
print(ml[MyEnum.ONE])
print(ml[1])  # error here

Why would you want to restrict access to a list to a different type than the default type for list index (i.e. int) though?

If you're looking to have specific attributes, why not define a class, or if you want to index through specific values, why not use a dict?

Perhaps you should provide some information on what you're actually hoping to do - Python is not C/C++ and it probably has some perfectly valid way of doing what you want to do, without having to emulate what C/C++ does instead.

If you're not looking to restrict, but just to iterate (which seems to be distinct from what the question stated):

for my_value in MyEnum:
    print(ml[my_value])

You could then leave out the assert statement, if you're not interesting in limiting access to indices with that type only.

Upvotes: 0

Nath
Nath

Reputation: 965

From the docs https://docs.python.org/3/library/enum.html Enumerations support iteration, in definition order so you can iterate over them to build lists, for example:

from enum import Enum

class Numbers(Enum):
    ONE = 1
    TWO = 2
    THREE = 3

assert [Number.value for Number in Numbers]  == [1,2,3]

Upvotes: 0

Related Questions