pythonpython-3.9

Is there a more elegant rewrite for this Python Enum value_of implementation?


I would like to get a value_of implementation for the StrEnum (Python 3.9.x). For example:

from enum import Enum


class StrEnum(str, Enum):
    """Enum with str values"""
    pass

class BaseStrEnum(StrEnum):
    """Base Enum"""
    @classmethod
    def value_of(cls, value): 
        try:
            return cls[value]
        except KeyError:
            try:
                return cls(value)
            except ValueError:
                return None

and then can use it like this:

class Fruits(BaseStrEnum):
    BANANA = "Banana"
    PEA = "Pea"
    APPLE = "Apple"

print(Fruits.value_of('BANANA'))
print(Fruits.value_of('Banana'))

it is just that the nested try-except doesn't look amazing, is there a better more idiomatic rewrite?


Solution

  • Since upon success of the first try block the function will return and won't execute the code that follows, there is no need to nest the second try block in the error handler of the first try block to begin with:

    def value_of(cls, value): 
        try:
            return cls[value]
        except KeyError:
            pass
        try:
            return cls(value)
        except ValueError:
            return None
    

    And since both of the error handlers are really meant to ignore the respective exceptions, you can use contextlib.suppress to simply suppress those errors:

    from contextlib import suppress
    
    def value_of(cls, value):
        with suppress(KeyError):
            return cls[value]
        with suppress(ValueError):
            return cls(value)
        # return None
    

    Note that a function returns None by default so you don't have to explicitly return None as a fallback unless you want to make it perfectly clear.