CODEWITHSUNDEEP
pythonenums
Jim
Jim

Reputation: 574

Adding Attributes And Starting Value To Python enum.IntEnum

How can I define a Python enum class that somehow derives from int, has a custom starting value, and adds custom attributes? I know how to derive from int using enum.IntEnum and set the starting value,

Goo = enum.IntEnum("Goo", "MOO FOO LOO", start = 42)

and how to add custom attributes to the base enum type,

class Goo(enum.Enum):
    MOO = (42, "forty-two")
    FOO = (43, "forty-three")
    LOO = (44, "forty-four")

    def __init__(self, value, alias):
        self._value = value   #self.value gives AttributeError, as expected, grrr...
        self.alias = alias

but how do I do all three? I've tried all manner of __new__() too,

class Goo(int, enum.Enum):
    MOO = (42, "forty-two")
    FOO = (43, "forty-three")
    LOO = (44, "forty-four")

    def __new__(cls, value, alias):
        self = super().__new__(cls, value)
        self.alias = alias

        return self

mostly with odd errors,

TypeError: int() takes at most 2 arguments (3 given)

in this case. Thanks.

Jim

Upvotes: 5

Views: 1598

Answers (1)

Ethan Furman
Ethan Furman

Reputation: 69051

The behavior you want isn't already available in the stdlib Enum. You could use itertools.count() to get the numbers:

from enum import IntEnum
from itertools import count

_goo_count = count(42)
class Goo(IntEnum):
    #
    MOO = "forty-two"
    FOO = "forty-three"
    LOO = "forty-four"
    #
    def __new__(cls, alias):
        value = next(_goo_count)
        member = int.__new__(cls, value)
        member._value_ = value
        member.alias = alias
        return member

or you can use aenum:

from aenum import IntEnum

class Goo(IntEnum):
    _start_ = 42
    _init_ = 'alias'
    #
    MOO = "forty-two"
    FOO = "forty-three"
    LOO = "forty-four"

either way, in use it looks like:

>>> Goo.MOO
<Goo.MOO: 42>

>>> Goo.MOO.alias
'forty-two'

>>> Goo.MOO == 42
True

A couple stylistic notes:

  • you don't need parentheses to define a tuple

    • except for an empty tuple
    • or some place (like a function call), where the command would be interpreted as not a tuple

  • white space is good ;-)

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

Upvotes: 5

Related Questions