Surya Kanagasabapathy

Enums with Attributes in Python

2 Dec 2024

Overview

This is a way to create Enums which take more than one attribute per value

from enum import Enum

class Severity(Enum):
	# Add all the attributes needed here
    string: str
    score: float

    def __new__(cls, string:str, score:float):
        obj = object.__new__(cls)
        value = len(cls.__members__) + 1
        # We can also assign any one of our attributes to __value__ and they will accessible by using the built-in .value
        # obj._value_ = value
        
        obj.string = string
        obj.score = score
        return obj

    LOW = "low",3.0
    MEDIUM = "medium",5.0
    HIGH = "high", 7.0

The earlier snippet is not useful when you don’t know the attributes that you will have. In that case, this will help but this will not have type hints for the attributes.

from enum import Enum

class DynamicEnumWithAttrs(Enum):
    def __new__(cls, kwargs):
        obj = object.__new__(cls)
        # Set the first attribute as the value if there are any attributes
        first_key = next(iter(kwargs), None)
        obj._value_ = kwargs[first_key] if first_key else None
        # Dynamically assign attributes
        for key, value in kwargs.items():
            setattr(obj, key, value)
        return obj

class Severity(DynamicEnumWithAttrs):
    # Define enum members with arbitrary attributes
    LOW = {"string": "low", "score": 3.0}
    MEDIUM = {"string": "medium", "score": 5.0}
    HIGH = {"string": "low", "score": 7.0, "test":"aa"}

Source

Tags