How can I extract a python enum subset without redefining it?

user_7535684 0
Viewed 2276

How can I extract a python enum subset without redefining it?

from enum import unique, Enum
@unique
class MyEnum(Enum):
    ONE = 1
    TWO = 2
    THREE = 3
    FOUR = 4

I would like to get an equivalent to MyDesiredSubset without having to define it again.

@unique
class MyDesiredSubset(Enum):
    THREE = 3
    FOUR = 4

So far I tried something like this, but MyTrySubset is broken and the code is ugly.

@unique
class MyTrySubset(Enum):
    pass

for item in MyEnum:
    setattr(MyTrySubset, item.name, item.value)

Any suggestions how to get MyDesiredSubset without redefining it?

enums python-3.x python
3 Answers

Your code won't work - enums are not quite classes:

@unique
class MyEnum(Enum):
    ONE = 1
    TWO = 2
    THREE = 3
    FOUR = 4

@unique
class MyTrySubset(Enum):
    pass

 for item in MyEnum: 
    setattr(MyTrySubset, item.name,     item.value)  # no duplication error by @unique
    setattr(MyTrySubset, item.name+"a", item.value)  # no duplication error by @unique

for s in MyTrySubset:
    print(s)           # no output - at all

for s in MyEnum:
    print(s)           # prints all repr() of all Enum-values defined

Use the other enum to declare this one (won't compare equal though):

@unique
class MyDesiredSubset(Enum):
    THREE = MyEnum.THREE
    FOUR = MyEnum.FOUR

or use the fluent approach:

MyOther = Enum("MyOther", [(a.name,a.value) for a in MyEnum 
                           if a in [MyEnum.THREE,MyEnum.FOUR]] )

If you use IntEnum instead, you can even compare them:

@unique
class MyIntEnum(IntEnum):
    ONE = 1
    TWO = 2
    THREE = 3
    FOUR = 4

@unique
class MyDesiredIntSubset(IntEnum):
    THREE = MyIntEnum.THREE
    FOUR = MyIntEnum.FOUR

print(MyDesiredSubset.THREE == MyEnum.THREE)       # False
print(MyDesiredIntSubset.THREE == MyIntEnum.THREE) # True 
print(MyDesiredIntSubset.THREE == 3)               # True @Steven Rumbalski
user_75053950

Patrick Artner's answer contains what I was looking for looping through an enum with an if statement keeping desired elements.

Just extracted the relevant bits.

MySubset = Enum("MySubet", [(a.name, a.value) for a in MyEnum if a.value > 2 ] )

Seems to be equivalent to (without the unique decorator):

class MySubset(Enum):
    THREE = 3
    FOUR = 4
user_75356840

I would suggest defining properties as indicated in this question:

from enum import unique, Enum
@unique
class MyEnum(Enum):
    ONE = 1
    TWO = 2
    THREE = 3
    FOUR = 4

    @property
    def my_desired_subset(self):
        return self.THREE, self.FOUR

    @property
    def is_desired_subset(self):
        return self in self.my_desired_subset

And now you can do:

MyEnum.my_desired_subset
user_71978990
Related