How to create Python Enum which allows synonyms?

Viewed 174

I'm normalising existing, messy data, and I'd like to create an Enum which allows synonyms for the canonical names of the members, so that if someone uses a synonym value when instantiating the enum, they will get the canonical one back. Ie.

class TrainOutcome(enum.Enum):
    PASSED = "PASSED"
    SUCCESS = "PASSED" # Deprecated synonym for "PASSED"
    FAILED = "FAILED"
    STARTED = "STARTED"

This executes fine, but the resulting enum doesn't behave as expected:

>>> TrainOutcome("PASSED")
<TrainOutcome.PASSED: 'PASSED'>

# I want to get <TrainOutcome.PASSED: 'PASSED'> here as well
>>> TrainOutcome("SUCCESS")
ValueError: 'SUCCESS' is not a valid TrainOutcome

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3.8/enum.py", line 309, in __call__
    return cls.__new__(cls, value)
  File "/usr/lib/python3.8/enum.py", line 600, in __new__
    raise exc
  File "/usr/lib/python3.8/enum.py", line 584, in __new__
    result = cls._missing_(value)
  File "/usr/lib/python3.8/enum.py", line 613, in _missing_
    raise ValueError("%r is not a valid %s" % (value, cls.__name__))
ValueError: 'SUCCESS' is not a valid TrainOutcome

This is despite the fact that the __members__ attribute seems to map things exactly the way I was hoping for:

>>> TrainOutcome.__members__
mappingproxy({'PASSED': <TrainOutcome.PASSED: 'PASSED'>, 'SUCCESS': <TrainOutcome.PASSED: 'PASSED'>, 'FAILED': <TrainOutcome.FAILED: 'FAILED'>, 'STARTED': <TrainOutcome.STARTED: 'STARTED'>})
>>> TrainOutcome['SUCCESS']
<TrainOutcome.PASSED: 'PASSED'>
>>> TrainOutcome['PASSED']
<TrainOutcome.PASSED: 'PASSED'>

How do I create the enum so that the constructor accepts and returns the same value as indexing the type does?

Edit: The existing Python Enum with duplicate values doesn't answer my question, since in essence it's trying to achieve the opposite of what I'm after. The OP there wanted to make the resulting values more distinct, I want to make them less distinct. In fact, the ideal solution would be not to have the synonym member at all (since I'm using the resulting Enum in SQLAlchemy context, which looks at the member names, not their values), and just silently replace "SUCCESS" with "PASSED" during construction time, but defining a custom __init__ on enums that calls super() doesn't seem to work.

Edit: This question and answer provides the easiest solution so far: use aenum.MultiValueEnum.

Otherwise, here's a homegrown solution that seems to be in the spirit of how you're supposed to do in Python 3.6+, somewhat inspired by @Green Cloak Guy's answer:

class EnumSynonymMixin:
    """
    Enum mixin which provides the ability to define synonyms,
    ie. values which can be passed into an enum's constructor, that
    name the same member as one of the defined values, without adding
    any extra members (useful for using with SQLAlchemy's Enum mapping)

    For example:

    class MyEnum(EnumSynonymMixin, enum.Enum):
        FOO = "FOO"
        BAR = "BAR"

        @classmethod
        def synonyms(cls):
            return {"MYFOO": "FOO"}
    
    >>> MyEnum("MYFOO")
    <MyEnum.FOO: 'FOO'>
    """
    @classmethod
    def synonyms(cls):
        """Override to provide a dictionary of synonyms for values that can be
        passed to the constructor"""
        return {}

    @classmethod
    def _missing_(cls, val):
        synonyms = cls.synonyms()
        if val in synonyms:
            return cls.__members__[synonyms[val]]
        return super()._missing(val)


class TrainOutcome(EnumSynonymMixin, enum.Enum):
    PASSED = "PASSED"
    FAILED = "FAILED"
    STARTED = "STARTED"

    @classmethod
    def synonyms(cls):
        return {"SUCCESS": "PASSED"}
1 Answers

This should do what you want. Essentially, the class wrapper makes TrainOutcome(value) behave like TrainOutcome[value] when the former would otherwise produce an error (as in the case you've described, where you're trying to call it with "SUCCESS"). It does this by intercepting the call to __new__() and replacing the first argument.

Per comments on your question, you probably shouldn't do this - there's little reason I can think of why TrainOutcome['SUCCESS'] shouldn't suffice for your needs.

def callActsLikeGetitem(c):
    oldnew = c.__new__
    def newwrapper(cls, *args, **kwargs):
        try:
            return oldnew(cls, *args, **kwargs)
        except ValueError:
            if len(args) > 0:
                args = (cls[args[0]].name, *args[1:])
            return oldnew(cls, *args, **kwargs)
    c.__new__ = newwrapper
    return c

@callActsLikeGetitem
class TrainOutcome(enum.Enum):
    PASSED = "PASSED"
    SUCCESS = "PASSED" # Deprecated synonym for "PASSED"
    FAILED = "FAILED"
    STARTED = "STARTED"

TrainOutcome("SUCCESS")
# <TrainOutcome.PASSED: 'PASSED'>
Related