35

I've been messing around with python's enum library and have come across a conundrum. In the docs, they show an example of an auto-numbering enum, wherein something is defined:

class Color(AutoNumber):
    red = ()
    green = ()
    ...

I want to make a similar class, but the value would automatically be set from the name of the member AND keep the functionality that you get from doing the str and enum mixin stuff

So something like:

class Animal(MagicStrEnum):
    horse = ()
    dog = ()

Animal.dog == 'dog' # True

I've looked at the source code of the enum module and tried a lot of variations messing around with __new__ and the EnumMeta class

Keelan Armstrong
  • 351
  • 1
  • 3
  • 3
  • 4
    do you mean `Animal.dog.value == 'dog'`? – Padraic Cunningham Aug 25 '15 at 21:57
  • related: http://www.acooke.org/cute/Pythonssad0.html – wim Aug 25 '15 at 22:27
  • I think it could be done by modifying EnumMeta._create_() in enum.py (https://hg.python.org/cpython/file/3.4/Lib/enum.py#l295, however that cannot be immediately overriden due to ValueError on sundered names coming from _is_sunder() to protect the code. However, if names = ['red','green','blue'] and Color = Enum('Color', names=zip(names,names)), then Color.red.value == 'red', Color.green.value == 'green' and Color.blue.value == 'blue'. –  Aug 25 '15 at 23:32
  • What are you trying to do here? – rlbond Aug 25 '15 at 23:53

2 Answers2

41

Update: 2017-03-01

In Python 3.6 (and Aenum 2.01) Flag and IntFlag classes have been added; part of that was a new auto() helper that makes this trivially easy:

>>> class AutoName(Enum):
...     def _generate_next_value_(name, start, count, last_values):
...         return name
...
>>> class Ordinal(AutoName):
...     NORTH = auto()
...     SOUTH = auto()
...     EAST = auto()
...     WEST = auto()
...
>>> list(Ordinal)
[<Ordinal.NORTH: 'NORTH'>, <Ordinal.SOUTH: 'SOUTH'>, <Ordinal.EAST: 'EAST'>, <Ordinal.WEST: 'WEST'>]

Original answer

The difficulty with an AutoStr class is that the name of the enum member is not passed into the code that creates it, so it is unavailable for use. Another wrinkle is that str is immutable, so we can't change those types of enums after they have been created (by using a class decorator, for example).

The easiest thing to do is use the Functional API:

Animal = Enum('Animal', [(a, a) for a in ('horse', 'dog')], type=str)

which gives us:

>>> list(Animal)
[<Animal.horse: 'horse'>, <Animal.dog: 'dog'>]

>>> Animal.dog == 'dog'
True

The next easiest thing to do, assuming you want to make a base class for your future enumeration use, would be something like my DocEnem:

class DocEnum(Enum):
    """
    compares equal to all cased versions of its name
    accepts a doctring for each member
    """
    def __new__(cls, *args):
        """Ignores arguments (will be handled in __init__)"""
        obj = object.__new__(cls)
        obj._value_ = None
        return obj

    def __init__(self, doc=None):
        # first, fix _value_
        self._value_ = self._name_.lower()
        self.__doc__ = doc

    def __eq__(self, other):
        if isinstance(other, basestring):
            return self._value_ == other.lower()
        elif not isinstance(other, self.__class__):
            return NotImplemented
        return self is other

    def __hash__(self):
        # keep DocEnum hashable
        return hash(self._value_)

    def __ne__(self, other):
        return not self == other

and in use:

class SpecKind(DocEnum):
    REQUIRED = "required value"
    OPTION = "single value per name"
    MULTI = "multiple values per name (list form)"
    FLAG = "boolean value per name"
    KEYWORD = 'unknown options'

Note that unlike the first option, DocEnum members are not strs.


If you want to do it the hard way: subclass EnumMeta and fiddle with the new Enum's class dictionary before the members are created:

from enum import EnumMeta, Enum, _EnumDict

class StrEnumMeta(EnumMeta):
    def __new__(metacls, cls, bases, oldclassdict):
        """
        Scan through `oldclassdict` and convert any value that is a plain tuple
        into a `str` of the name instead
        """
        newclassdict = _EnumDict()
        for k, v in oldclassdict.items():
            if v == ():
                v = k
            newclassdict[k] = v
        return super().__new__(metacls, cls, bases, newclassdict)

class AutoStrEnum(str, Enum, metaclass=StrEnumMeta):
    "base class for name=value str enums"

class Animal(AutoStrEnum):
    horse = ()
    dog = ()
    whale = ()

print(Animal.horse)
print(Animal.horse == 'horse')
print(Animal.horse.name, Animal.horse.value)

Which gives us:

Animal.horse
True
horse horse

1 Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

Ethan Furman
  • 63,992
  • 20
  • 159
  • 237
  • @VillasV: Thanks for finding that error! I'm sorry the reviewers didn't see it was a correct edit. – Ethan Furman Nov 01 '16 at 19:30
  • Brilliant answer! Just one problem: Although subclassed from `Enum`, `DocEnum` seems to not have a `__hash__` method. I have to manually copy `Enum.__hash__` to make `DocEnum` hashable. Why is that? – ComeOnGetMe Jun 14 '19 at 18:37
  • @ComeOnGetMe: In Python 3 if an `__eq__` is defined then `__hash__` is set to `None` unless `__hash__` is defined (or set to something else) in the class. Answer updated. – Ethan Furman Jun 18 '19 at 16:11
  • This no longer works in 3.9 because the private `_EnumDict` has changed. – gps Nov 11 '21 at 00:16
  • @gps: I just tested the top most solution of my answer, and the original solution in my answer, from Pythons 3.6 - 3.11, and both worked in every version. Exactly what did you try? – Ethan Furman Nov 11 '21 at 06:17
  • The latter one, the subclass, doesn't work as of 3.9. – gps Nov 12 '21 at 16:34
  • @gps (forgot to mention you in my reply): As I said yesterday, I have copy/pasted that code into a file, ran it with Python 3.9 (in fact, all Pythons between 3.6 and 3.11, inclusive), and I get the correct results with 3.9 and all the others. (And I just rebuilt 3.9.8 to make sure.) Can you email me the exact details of what you are doing? If you've found a bug I'd like to fix it. Oh, and the only change to `_EnumDict` in 3.9 is a warning if the member name is name mangled. – Ethan Furman Nov 16 '21 at 16:15
1

Perhaps you are looking for the name attribute which is automatically provided by the Enum class

>>> class Animal(Enum):
...     ant = 1
...     bee = 2
...     cat = 3
...     dog = 4
...

>>> Animal.ant.name == "ant"
True

Though if you really want to shoot yourself in the foot. And I'm sure this will introduce a whole world of gotchas (I've eliminated the most obvious one).

from enum import Enum, EnumMeta, _EnumDict

class AutoStrEnumDict(_EnumDict):
    def __setitem__(self, key, value):
        super().__setitem__(key, key)

class AutoStrEnumMeta(EnumMeta):
    @classmethod
    def __prepare__(metacls, cls, bases):
        return AutoStrEnumDict()
    def __init__(self, name, bases, attrs):
        super().__init__(name, bases, attrs)
        # override Enum.__str__
        # can't put these on the class directly otherwise EnumMeta overwrites them
        # should also consider resetting __repr__, __format__ and __reduce_ex__
        if self.__str__ is not str.__str__:
            self.__str__ = str.__str__

class AutoStrNameEnum(str, Enum, metaclass=AutoStrEnumMeta):
    pass

class Animal(AutoStrNameEnum):
    horse = ()
    dog = ()

print(Animal.horse)
assert Animal.horse == "horse"
assert str(Animal.horse) == "horse" 
# and not equal to "Animal.horse" (the gotcha mentioned earlier)
Dunes
  • 37,291
  • 7
  • 81
  • 97
  • Your modification to `_EnumDict` makes it impossible to add methods as they would also be converted to `str`s. – Ethan Furman Aug 31 '15 at 14:45
  • 1
    Easy to get around. Either check to see if the value is a function. Or to make it more robust, create a sentinel value and only do the conversion when the value is the sentinel. – Dunes Aug 31 '15 at 14:52