Python - extend enum fields during creation

129 Views Asked by At

Is it possible extend enum during creation? Example:

class MyEnum(enum.StrEnum):
    ID = "id"
    NAME = "NAME

And I need that after creating this enum contains the next fields:

ID = "id"
NAME = "name"
ID_DESC = "-id"
NAME_DESC = "-name"

I need this to create custom ordering enum for FastAPI project

Now I have the next way to create new enum

NewEnum = enum.StrEnum(
    f"{name.title()}OrderingEnum",
    [
        (
            f"{ordering_field.upper()}_DESC"
            if ordering_field.startswith("-")
            else ordering_field.upper(),
            ordering_field,
        )
        for ordering_field in itertools.chain(
            values,
            [f"-{field}" for field in values],
        )
    ],
)

But I need do this automatically, because each module with model have similar enum. Maybe this is possible to solve my problem using MetaClass for my enum class, or override __new__ method, but I didn't find working solution yet

1

There are 1 best solutions below

1
Ethan Furman On BEST ANSWER

Enhancing the __new__ of EnumType (used to be EnumMeta and still has that as an alias) will do the job:

from enum import EnumMeta, StrEnum

class IDEnumMeta(EnumMeta):
    def __new__(metacls, cls, bases, classdict, **kwds):
        # add new entries to classdict
        for name in list(classdict._member_names):
            classdict[f'{name}_DESC'] = f'-{classdict[name]}'
        return super().__new__(metacls, cls, bases, classdict, **kwds)

class IDEnum(StrEnum, metaclass=IDEnumMeta):
    pass

and in use:

>>> class MyEnum(IDEnum):
...     ID = 'id'
...     NAME = 'name'

>>> list(MyEnum)
[<MyEnum.ID: 'id'>, <MyEnum.NAME: 'name'>, <MyEnum.ID_DESC: '-id'>, <MyEnum.NAME_DESC: '-name'>]

This answer currently uses the private _EnumDict and its _member_names attribute, but they'll be public in 3.13.


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