Validate value is in Python Enum values

user_6220759 0
Viewed 6295

I've inherited code that looks something like this.

class Clients(IntEnum):
    ALICE = 1
    BOB = 2
    PETER = 3
    CHERYL = 4
    LARRY = 5

if client_id == 1:
    client_info = find_info(Clients.ALICE.value)
elif client_id == 2:
    client_info = find_info(Clients.BOB.value)
elif client_id == 3:
    client_info = find_info(Clients.PETER.value)
elif client_id == 4:
    client_info = find_info(Clients.CHERYL.value)
elif client_id == 5:
    client_info = find_info(Clients.LARRY.value)
else:
    raise Exception('Unknown client_id.')

Not having much experience with Python enum, I have a burning desire to simplify this into something like this (pseudocode):

if client_id in dict(Clients).keys():
    client_info = find_info(client_id)
else:
    raise Exception('Unknown client_id.')

I've tried Clients.__members__ and Clients.__dict__, but they don't quite behave as I'd expect, returning something called a mappingproxy.

I can do 'ALICE' in Clients.__members__, but is there an equivalent for the values? If not, is there a more elegant way to write this bit of code?

enums python
3 Answers

You can store only values in a plain tuple (or list) and use the fact that enums allow access by value:

values = tuple(item.value for item in Clients)

if client_id in values:
    # do something with Clients(client_id)
else:
    # print message or raise exception

Or you can map values to enum members using a dictionary:

mapping = dict((item.value, item) for item in Clients)

if client_id in mapping:
    # do something with mapping[client_id]
else:
    # print message or raise exception
user_87950730 
try:
    client_info = find_info(Clients(client_id))
except ValueError:
    # customize error message
    raise Exception('Unknown client_id.')
user_2088800

Here is how you can validate values with helper functions:

For the Enums:

class Customers(IntEnum):
    ABC = 1
    FOO = 2
    JDI = 3

class Clients(IntEnum):
    ALICE = 1
    BOB = 2
    PETER = 3

Implement helper functions:

from enum import IntEnum, EnumMeta
def isEnumClass(o) -> bool:
    return type(o) == EnumMeta

def isEnumMember(o) -> bool:
    return type(o).__class__ == EnumMeta

def enumHasMember(en, o, int_val:bool=False) -> bool:
    assert isEnumClass(en),\
        "param en is not an Enum class. Got %s!" \
            % ('(none)' if en is None else type(o).__name__)
    if type(o) == int:
        if not int_val: return False
        return len([m.value for m in en if m.value == o]) == 1
    else:
        return not o is None and o.__class__ == en

Usage:


print("1: %s" % enumHasMember(Customers, Customers.ABC))
print("2: %s" % enumHasMember(Customers, 1, int_val=False))
print("3: %s" % enumHasMember(Customers, 1, int_val=True))
print("4: %s" % enumHasMember(Customers, 4, int_val=True))
print("5: %s" % enumHasMember(Customers, Clients.ALICE))

> 1: True
> 2: False
> 3: True
> 4: False
> 5: False

Alternatively, if you control all of the code, you can create create a classmethod for a custom IntEnu class:

from enum import IntEnum, EnumMeta
class MyIntEnum(IntEnum):
    @classmethod
    def hasMember(cls, o, strict:bool=True) -> bool:
        if type(o) == int:
            if strict: return False
            return len([m.value for m in cls if m.value == o]) == 1
        else:
            return not o is None and o.__class__ == cls


class Clients(MyIntEnum):
    ALICE = 1
    BOB = 2
    PETER = 3
    CHERYL = 4
    LARRY = 5
    
class Customers(MyIntEnum):
    ABC = 1
    FOO = 2
    JDI = 3

Usage:

print("1: %s" % Customers.hasMember(Customers.ABC))
print("2: %s" % Customers.hasMember(1))
print("3: %s" % Customers.hasMember(1, strict=False))
print("4: %s" % Customers.hasMember(4, strict=False))
print("5: %s" % Customers.hasMember(Clients.ALICE))

> 1: True
> 2: False
> 3: True
> 4: False
> 5: False

Another handy class method would be validate method for a one line hard assertion:

class MyIntEnum(IntEnum):
        ...
    @classmethod
    def validate(cls, alias:str, o, strict:bool=True):
        assertNotBlank('alias', alias)
        assert cls.hasMember(o, strict),\
            f"Argument '{alias}' is not a member of MyIntEnum {cls.__module__}.{type(cls).__name__}. Got: {dump(o, True)}"
        
        return o

Please let me know if there is any issue with the above design. Just started researching this myself.

user_2868070
Related