Reputation: 100196
python enums with attributes
Consider:
class Item:
def __init__(self, a, b):
self.a = a
self.b = b
class Items:
GREEN = Item('a', 'b')
BLUE = Item('c', 'd')
Is there a way to adapt the ideas for simple enums to this case? (see this question) Ideally, as in Java, I would like to cram it all into one class.
Java model:
enum EnumWithAttrs {
GREEN("a", "b"),
BLUE("c", "d");
EnumWithAttrs(String a, String b) {
this.a = a;
this.b = b;
}
private String a;
private String b;
/* accessors and other java noise */
}
Upvotes: 85
Views: 75478
Answers (11)
Reputation: 1
This might be a bit overkill, but here is a class that has a little more functionality and customizability and (most importantly) will not generate any mypy errors.
from enum import Enum
from typing import Any, List
class MultiAttributeEnum(Enum):
def __init__(self, attributes: Any) -> None:
"""
A subclass of Enum that allows members to have multiple attributes.
Members can be defined by either dictionaries or lists. The attributes
of the members must also be defined in the subclass to avoid mypy issues
Example Usage:
Animals(MultiAttributeEnum):
# Define Attributes
species: str
mammal: bool
height: float
# Define Members
Dog = {'species': 'Canis Lupus Familiaris', 'mammal': True, 'height': 0.5}
Giraffe = {'species': 'Giraffa', 'mammal': True, 'height': 10}
By default the main attribute is the first value in the dictionary. So
for the above example:
>>> Animals('Canis Lupus Familiaris')
Animals.Dog
The main attribute can be changed by overriding
self.get_main_attribute. If the value of the main attribute is not
unique to a specific member, then the member that was defined first is
returned.
Members can also be defined using lists so long as the
self.get_attribute_keys() method is also defined by the user.
Additionally, not all members need to have all the attributes. This is
also true when the members are defined using dictionaries. You can even
mix dictionaries and lists.
Example Usage:
Animals(MultiAttributeEnum):
# Define Attributes
species: str
mammal: bool
height: int
# Specify attribute keys so we can define members with list
def get_attribute_keys(self, attributes):
return ['species', 'mammal', 'height']
# Define Members
Dog = ['Canis Lupus Familiaris', True] # Dog will be missing the height attribute
Giraffe = {'species': 'Giraffa', 'mammal': True, 'height': 10}
Cat = ['Felis Catus', True, 0.25]
Args:
attributes (Union[Dict[str, Any], List]): A list or dictionary of
attribute values for the enum member. If a list is given then
self.get_attribute_keys must be defined so that the class
knows what to name the attributes.
"""
if isinstance(attributes, dict):
for key, value in attributes.items():
setattr(self, key, value)
elif isinstance(attributes, list):
attribute_keys = self.get_attribute_keys(attributes)
for key, value in zip(attribute_keys, attributes):
setattr(self, key, value)
# Creat attributes that will be assigned for each member seperately
self.attribute_keys: List[str]
self.attribute_values: List[Any]
def __new__(cls, attributes: Any) -> Enum: # type: ignore
"""
Creates a new member. There is some Enum Weirdness here which is why we
have to tell mypy to ignore the typing. cls ends up being type
type[MultiAttributeEnum] but class methods expect type
MultiAttributeEnum for self. Additionally the return type is expected
to be MultiAttributeEnum but we can't specify that right now because
the class hasn't been instantiated. There might be a way around this
using bound TypeVars and generics but I don't know enough to figure it
out right now.
"""
# Create member
obj = object.__new__(cls)
# Set the attribute keys, values and main value for the member
if isinstance(attributes, dict):
obj._value_ = cls.get_main_attribute(cls, attributes) # type: ignore
obj.attribute_values = list(attributes.values())
obj.attribute_keys = list(attributes.keys())
if isinstance(attributes, list):
obj._value_ = cls.get_main_attribute(cls, attributes) # type: ignore
obj.attribute_values = attributes
obj.attribute_keys = cls.get_attribute_keys(cls, attributes)[: len(attributes)] # type: ignore
else: # Assume we only have a single attribute
obj.attribute_values = [attributes]
obj._value_ = attributes # reset the main value
# Return the member class
return obj # Return member
def keys(self) -> List[str]:
"""
Returns:
List[str]: a list containing the names of the attributes for this member
"""
return self.attribute_keys # These are set in __new__ for each member
def values(self) -> List[Any]:
"""
Returns:
List[Any]: A list of the attribute values for this member
"""
return self.attribute_values
def get_main_attribute(self, attributes: Any) -> Any:
"""
Returns the main attribute value given a members attributes. Whatever
is returned is what will be used to identify the member when calling
the class.
Example:
Animals(MultiAttributeEnum):
# Define Attributes
species: str
mammal: bool
height: int
# Define Members
Dog = {'species': 'Canis Lupus Familiaris', 'mammal': True, 'height': 0.5}
Giraffe = {'species': 'Giraffa', 'mammal': True, 'height': 10}
def (self, attributes):
return attributes['height']
>>> Animals(10)
Animals.Giraffe
>>> Animals(0.5)
Animals.Dog
Args:
attributes (Any): The attributes used to define the members
Returns:
Any: _description_
"""
if isinstance(attributes, dict):
return list(attributes.values())[0]
elif isinstance(attributes, list):
return attributes[0]
def get_attribute_keys(self, attributes: List) -> List[str]:
raise NotImplementedError(
"Received a list of attributes but the self.get_attribute_keys class method was not implemented"
)
Upvotes: 0
Reputation: 2497
These are two working examples:
from enum import Enum
class StatusInt(int, Enum):
READY = (0, "Ready to go!")
ERROR = (1, "Something wrong!")
def __new__(cls, value, description):
obj = int.__new__(cls, value)
obj._value_ = value
obj._description_ = description
return obj
@property
def description(self):
return self._description_
class StatusObj(Enum):
READY = (0, "Ready to go!")
ERROR = (1, "Something wrong!")
def __init__(self, value, description):
self._value_ = value
self._description_ = description
@property
def description(self):
return self._description_
print(str(StatusInt.READY == StatusInt.ERROR))
print(str(StatusInt.READY.value))
print(StatusInt.READY.description)
print(str(StatusObj.READY == StatusObj.ERROR))
print(str(StatusObj.READY.value))
print(StatusObj.READY.description)
Outputs:
False
0
Ready to go!
False
0
Ready to go!
References:
- https://docs.python.org/3/library/enum.html#when-to-use-new-vs-init
- https://docs.python.org/3/library/enum.html#planet
Upvotes: 11
Reputation: 11
extended-enum allows you to store additional information for each member of the enumeration.
An example with a color might look like this:
from dataclasses import dataclass
from extended_enum import ExtendedEnum, EnumField, BaseExtendedEnumValue
@dataclass(frozen=True)
class RGB:
red: int
green: int
blue: int
@dataclass(frozen=True)
class ColorExtendedEnumValue(BaseExtendedEnumValue):
rgb: RGB
hex: str
class Color(ExtendedEnum):
BLACK = EnumField(ColorExtendedEnumValue(value='black', rgb=RGB(0, 0, 0), hex='000000'))
GRAY = EnumField(ColorExtendedEnumValue(value='gray', rgb=RGB(128, 128, 128), hex='808080'))
PURPLE = EnumField(ColorExtendedEnumValue(value='purple', rgb=RGB(128, 0, 128), hex='800080'))
GREEN = EnumField(ColorExtendedEnumValue(value='green', rgb=RGB(0, 255, 0), hex='00FF00'))
BLUE = EnumField(ColorExtendedEnumValue(value='blue', rgb=RGB(0, 0, 255), hex='0000FF'))
>>> Color.PURPLE.value
'purple'
>>> Color.PURPLE.extended_value
ColorExtendedEnumValue(value='purple', rgb=RGB(red=128, green=0, blue=128), hex='800080')
>>> Color.PURPLE.extended_value.rgb == RGB(red=128, green=128, blue=128)
False
>>> Color.PURPLE.extended_value.rgb == RGB(red=128, green=0, blue=128)
True
>>> Color.PURPLE.extended_value.hex == '800080'
True
If you don't need to compare each field, you can use
field(compare=False)
from dataclasses import dataclass, field
from extended_enum import BaseExtendedEnumValue
@dataclass(frozen=True)
class ColorExtendedEnumValue(BaseExtendedEnumValue):
rgb: RGB = field(compare=False)
hex: str = field(compare=False)
In addition, a class is available out of the box that contains a description:
from extended_enum import ExtendedEnum, EnumField, ValueWithDescription
class Status(ExtendedEnum):
READY = EnumField(ValueWithDescription(value='ready', description='I am ready to do whatever is needed'))
ERROR = EnumField(ValueWithDescription(value='error', description='Something went wrong here'))
>>> Status.READY
<Status.READY: ValueWithDescription(value='ready', description='I am ready to do whatever is needed')>
>>> Status.READY.value
'ready'
>>> Status.READY.extended_value
ValueWithDescription(value='ready', description='I am ready to do whatever is needed')
>>> Status.READY.extended_value.description
'I am ready to do whatever is needed'
Upvotes: 0
Reputation: 1125018
Before Python 3.4 and the addition of the excellent enum module, a good choice would have been to use a namedtuple:
from collections import namedtuple
Item = namedtuple('abitem', ['a', 'b'])
class Items:
GREEN = Item('a', 'b')
BLUE = Item('c', 'd')
These days, any supported version of Python has enum, so please use that module. It gives you a lot more control over how each enum value is produced.
If you give each item a tuple of values, then these are passed to the __init__ method as separate (positional) arguments, which lets you set additional attributes on the enum value:
from enum import Enum
class Items(Enum):
GREEN = ('a', 'b')
BLUE = ('c', 'd')
def __init__(self, a, b):
self.a = a
self.b = b
This produces enum entries whose value is the tuple assigned to each name, as well as two attributes a and b:
>>> Items.GREEN, Items.BLUE
(<Items.GREEN: ('a', 'b')>, <Items.BLUE: ('c', 'd')>)
>>> Items.BLUE.a
'c'
>>> Items.BLUE.b
'd'
>>> Items(('a', 'b'))
<Items.GREEN: ('a', 'b')>
Note that you can look up each enum value by passing in the same tuple again.
If the first item should represent the value of each enum entry, use a __new__ method to set _value_:
from enum import Enum
class Items(Enum):
GREEN = ('a', 'b')
BLUE = ('c', 'd')
def __new__(cls, a, b):
entry = object.__new__(cls)
entry.a = entry._value_ = a # set the value, and the extra attribute
entry.b = b
return entry
def __repr__(self):
return f'<{type(self).__name__}.{self.name}: ({self.a!r}, {self.b!r})>'
I added a custom __repr__ as well, the default only includes self._value_. Now the value of each entry is defined by the first item in the tuple, and can be used to look up the enum entry:
>>> Items.GREEN, Items.BLUE
(<Items.GREEN: ('a', 'b')>, <Items.BLUE: ('c', 'd')>)
>>> Items.BLUE.a
'c'
>>> Items.BLUE.b
'd'
>>> Items('a')
<Items.GREEN: ('a', 'b')>
See the section on __init__ vs. __new__ in the documentation for further options.
Upvotes: 48
Reputation: 243
enum-properties provides an extension of the Enum base class that allows attributes on enum values and also allows symmetric mapping backwards from attribute values to their enumeration values.
Add properties to Python enumeration values with a simple declarative syntax. Enum Properties is a lightweight extension to Python's Enum class. Example:
from enum_properties import EnumProperties, p
from enum import auto
class Color(EnumProperties, p('rgb'), p('hex')):
# name value rgb hex
RED = auto(), (1, 0, 0), 'ff0000'
GREEN = auto(), (0, 1, 0), '00ff00'
BLUE = auto(), (0, 0, 1), '0000ff'
# the named p() values in the Enum's inheritance become properties on
# each value, matching the order in which they are specified
Color.RED.rgb == (1, 0, 0)
Color.GREEN.rgb == (0, 1, 0)
Color.BLUE.rgb == (0, 0, 1)
Color.RED.hex == 'ff0000'
Color.GREEN.hex == '00ff00'
Color.BLUE.hex == '0000ff'
Properties may also be symmetrically mapped to enumeration values, using s() values:
from enum_properties import EnumProperties, s
from enum import auto
class Color(EnumProperties, s('rgb'), s('hex', case_fold=True)):
RED = auto(), (1, 0, 0), 'ff0000'
GREEN = auto(), (0, 1, 0), '00ff00'
BLUE = auto(), (0, 0, 1), '0000ff'
# any named s() values in the Enum's inheritance become properties on
# each value, and the enumeration value may be instantiated from the
# property's value
Color((1, 0, 0)) == Color.RED
Color((0, 1, 0)) == Color.GREEN
Color((0, 0, 1)) == Color.BLUE
Color('ff0000') == Color.RED
Color('FF0000') == Color.RED # case_fold makes mapping case insensitive
Color('00ff00') == Color.GREEN
Color('00FF00') == Color.GREEN
Color('0000ff') == Color.BLUE
Color('0000FF') == Color.BLUE
Color.RED.hex == 'ff0000'
Upvotes: 3
Reputation: 2112
Here's another approach which I think is simpler than the others, but allows the most flexibility:
from collections import namedtuple
from enum import Enum
class Status(namedtuple('Status', 'name description'), Enum):
READY = 'ready', 'I am ready to do whatever is needed'
ERROR = 'error', 'Something went wrong here'
def __str__(self) -> str:
return self.name
It works as expected:
>>> str(Status.READY)
ready
>>> Status.READY
<Status.READY: Status(name='ready', description='I am ready to do whatever is needed')>
>>> Status.READY.description
'I am ready to do whatever is needed'
>>> Status.READY.value
Status(name='ready', description='I am ready to do whatever is needed')
Also you are able to retrieve the enum by name (Thanks @leoll2 for pointing this out). For example
>>> Status['READY']
<Status.READY: Status(name='ready', description='I am ready to do whatever is needed')>
You get the best of namedtuple and Enum.
Upvotes: 45
Reputation: 69288
Python 3.4 has a new Enum data type (which has been backported as enum34 and enhanced as aenum1). Both enum34 and aenum2 easily support your use case:
aenum(Python 2/3)import aenum class EnumWithAttrs(aenum.AutoNumberEnum): _init_ = 'a b' GREEN = 'a', 'b' BLUE = 'c', 'd'enum34(Python 2/3) or standard libraryenum(Python 3.4+)import enum class EnumWithAttrs(enum.Enum): def __new__(cls, *args, **kwds): value = len(cls.__members__) + 1 obj = object.__new__(cls) obj._value_ = value return obj def __init__(self, a, b): self.a = a self.b = b GREEN = 'a', 'b' BLUE = 'c', 'd'
And in use:
>>> EnumWithAttrs.BLUE
<EnumWithAttrs.BLUE: 1>
>>> EnumWithAttrs.BLUE.a
'c'
1 Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.
2 aenum also supports NamedConstants and metaclass-based NamedTuples.
Upvotes: 83
Reputation: 5461
For keyword-based initialization of attributes, you might try data-enum, a more lightweight implementation of enum with cleaner syntax for some cases, including this one.
from data_enum import DataEnum
class Item(DataEnum):
data_attribute_names = ('a', 'b')
Item.GREEN = Item(a='a', b='b')
Item.BLUE = Item(a='c', b='d')
I should note that I am the author of data-enum, and built it specifically to address this use case.
Upvotes: 3
Reputation: 15614
for small enums @property might work:
class WikiCfpEntry(Enum):
'''
possible supported storage modes
'''
EVENT = "Event"
SERIES = "Series"
@property
def urlPrefix(self):
baseUrl="http://www.wikicfp.com/cfp"
if self==WikiCfpEntry.EVENT:
url= f"{baseUrl}/servlet/event.showcfp?eventid="
elif self==WikiCfpEntry.SERIES:
url= f"{baseUrl}/program?id="
return url
Upvotes: 4
Reputation: 1376
Inspired by some of the other answers, I found a way of including additional fields to an enum as 'transparently' as possible, overcoming some shortcomings of the other approaches. Everything works the same as if the additional fields weren't there.
The enum is immutable just like a tuple, the value of the enum is just as it would be without the additional fields, it works just like a normal enum with auto(), and selecting an enum by value works.
import enum
# Common base class for all enums you want to create with additional fields (you only need this once)
class EnumFI(enum.Enum):
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
cls._values = []
def __new__(cls, *args, **kwargs):
value = args[0]
if isinstance(value, enum.auto):
if value.value == enum._auto_null:
value.value = cls._generate_next_value_(None, 1, len(cls.__members__), cls._values[:]) # Note: This just passes None for the key, which is generally okay
value = value.value
args = (value,) + args[1:]
cls._values.append(value)
instance = cls._member_type_.__new__(cls, *args, **kwargs)
instance._value_ = value
return instance
def __format__(self, format_spec):
return str.__format__(str(self), format_spec)
Then anywhere in the code you can just do:
from enum import auto
from collections import namedtuple
class Color(namedtuple('ColorTuple', 'id r g b'), EnumFI):
GREEN = auto(), 0, 255, 0
BLUE = auto(), 0, 0, 255
Example output:
In[4]: Color.GREEN
Out[4]: <Color.GREEN: 1>
In[5]: Color.GREEN.value
Out[5]: 1
In[6]: Color.GREEN.r
Out[6]: 0
In[7]: Color.GREEN.g
Out[7]: 255
In[8]: Color.GREEN.b
Out[8]: 0
In[9]: Color.GREEN.r = 8
Traceback (most recent call last):
File "/home/phil/anaconda3/envs/dl/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3326, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-9-914a059d9d3b>", line 1, in <module>
Color.GREEN.r = 8
AttributeError: can't set attribute
In[10]: Color(2)
Out[10]: <Color.BLUE: 2>
In[11]: Color['BLUE']
Out[11]: <Color.BLUE: 2>
Upvotes: 2
Reputation: 661
For Python 3:
class Status(Enum):
READY = "ready", "I'm ready to do whatever is needed"
ERROR = "error", "Something went wrong here"
def __new__(cls, *args, **kwds):
obj = object.__new__(cls)
obj._value_ = args[0]
return obj
# ignore the first param since it's already set by __new__
def __init__(self, _: str, description: str = None):
self._description_ = description
def __str__(self):
return self.value
# this makes sure that the description is read-only
@property
def description(self):
return self._description_
And you can use it as a standard enum or factory by type:
print(Status.READY)
# ready
print(Status.READY.description)
# I'm ready to do whatever is needed
print(Status("ready")) # this does not create a new object
# ready
Upvotes: 56
Related Questions
- Python Enum with multiple attributes
- Enums with Associated Values in Python
- Enum of enums in Python?
- attributes of attributes of Enums?
- Python metaclass for Enum members with variable number of attributes
- Enum having multiple attributes as a Collection Constant
- Using Enums to inspect attributes of a class/instance in python
- Defining python classes with enums
- Enum with attributes in Python
- python: using attributes as an enum, and then assigning the attributes to another class