How to add extra data to TextChoices?

Viewed 314

How Can I add extra data to django.db.models.TextChoices?

class Fruit(models.TextChoices):

    APPLE = ('myvalue', True, 'mylabel')

such that:

>>> Fruit.APPLE.is_tasty
True
>>> # And it still works otherwise
>>> Fruit.APPLE.value
'myvalue'
>>> Fruit.APPLE.label
'mylabel'
1 Answers

You need to do something similar to what the python Enum docs suggest, but unlike python Enum, the label is already taken care of by models.Choices:

class Fruit(models.TextChoices):

    APPLE = ('myvalue', True, 'mylabel')

    def __new__(cls, value, is_tasty):
        obj = str.__new__(cls, value)
        obj._value_ = value
        obj.is_tasty = is_tasty
        return obj

If you use this on an IntegerChoices, you'll need int.__new__. If you use __init__ instead of __new__, the enum value would become ('myvalue', True), which gets used in Fruit.choices and probably won't fit your model Field.

Though note when the choices enum is used on a model field, you never actually pass it the enum, so it's not aware of the enum. E.g. the form field derived from it with ModelForm will treat them as str values and after a POST the field value will be a regular str, not an enum value. For the form case you can define MyForm.clean_fruit or provide a manual form field TypedChoiceField(coerce=Fruit), elsewhere you may need to look up the enum value again with Fruit(value) or you could add this mixin to your field:

class EnumMixin:

    'Convert a DB value back to its Choices value'

    def __init__(self, *args, enum: models.Choices, **kwargs):
        self.__enum = enum
        # it sets choices for you using the enum
        super().__init__(*args, choices=enum.choices, **kwargs)

    def deconstruct(self):
        'Get constructor args to reconstruct this field with later'
        name, path, args, kwargs = super().deconstruct()
        kwargs['enum'] = self.__enum
        del kwargs['choices']
        return name, path, args, kwargs

    def from_db_value(self, value, expression, connection):
        # Convert from db value
        return self.__to_enum(value)

    def to_python(self, value):
        'Called by deserialization and during clean() method used in forms'
        return self.__to_enum(value)

    def __to_enum(self, value):
        if value is None:
            return None
        return self.__enum(value)

class EnumCharField(EnumMixin, models.CharField):
    pass

class MyModel(models.Model):
    field = EnumCharField(enum=Fruit, ...)

deconstruct is used by django migrations though note it won't reconstruct the Enum at the time you make the migration, it will use whatever your Enum is at the time you apply the migration.

Why not call super?

While super().__new__ works fine in general, this is not the case for Enum subclasses (models.TextChoices is a models.Choices which is an Enum). Python docs note the following:

The __new__() method, if defined, is used during creation of the Enum members; it is then replaced by Enum’s __new__() which is used after class creation for lookup of existing members.

So, EnumMeta replaces the class' __new__ with Enum.__new__ on Fruit, TextChoices, ... If you call super().__new__ in Fruit.__new__ then this calls TextChoices.__new__ which is actually Enum.__new__ and which will not expect the args you are passing it (and even if it did accept your args it never calls super().__new__ itself).

It would raise:

ValueError: 'myvalue' is not a valid Fruit
Related