Django Custom EnumField

enumfield.py

from enum import Enum
from typing import List, Optional, Tuple, Type, Union

from django.contrib.auth.models import AbstractUser
from django.db import models
from django.db.backends.base.base import BaseDatabaseWrapper
from django.db.models.expressions import Expression
from django.forms import TypedChoiceField


class EnumFormField(TypedChoiceField):
    def prepare_value(self, value: Union[Enum, str]):  # for django.forms.boundfield
        if isinstance(value, Enum):
            return value.name
        return value


class EnumField(models.CharField):
    def __init__(self, enum: Type[Enum], *args, **kwargs):
        self.enum = enum
        names = [item.name for item in enum]
        kwargs["max_length"] = max([len(name) for name in names])
        kwargs["default"] = kwargs.get("default", self.enum[names[0]])
        kwargs["choices"] = [(item.name, item.value) for item in enum]
        kwargs["blank"] = False
        kwargs["null"] = False
        super().__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args_, kwargs = super().deconstruct()
        del kwargs["max_length"]
        del kwargs["default"]
        del kwargs["choices"]
        return name, path, args_ + [self.enum], kwargs

    def get_prep_value(self, value: Enum):
        return value.name

    def from_db_value(
        self, value: str, expression: Expression, connection: BaseDatabaseWrapper
    ):  # pylint: disable=unused-argument
        return self.enum[value]

    def to_python(self, value: str) -> Optional[Enum]:
        if isinstance(value, self.enum):
            return value

        if value is None:
            return value

        return self.enum[value]

    def formfield(self, **kwargs):
        kwargs["choices_form_class"] = EnumFormField
        return super().formfield(**kwargs)

    def clean(self, value: str, model_instance: models.Model) -> Optional[Enum]:
        enum_val = self.to_python(value)
        self.validate(enum_val and enum_val.name, model_instance)
        self.run_validators(enum_val and enum_val.name)
        return enum_val

    @property
    def flatchoices(self) -> List[Tuple[Enum, str]]:
        """
        Allows for read-only support.
        Overrides django.db.models.fields.Field.flatchoices.
        django.contrib.admin.helpers.AdminReadonlyField.contents calls display_for_fields
        which will attempt to get the display value through flatchoices. This will fail if
        the field value portion of the items within the choices tuple is a string rather
        than an enum value.
        """
        return [(self.enum[choice], value) for choice, value in self.choices]


# Usage in model


@unique
class FooChoice(Enum):  # must be at top-level for generating migrations properly
    CHOICE_1 = "Choice 1"
    CHOICE_2 = "CHoice 2"

    
class Foo(models.Model):
    Choice = FooChoice
    
    choice: FooChoice = EnumField(FooChoice)