Custom Django field does not return Enum instances from query

979 views Asked by At

I have a simple custom field implemented to utilize Python 3 Enum instances. Assigning enum instances to my model attribute, and saving to the database works correctly. However, fetching model instances using a QuerySet results in the enum attribute being a string, instead of the respective Enum instance.

How do I get the below EnumField to return valid Enum instances, rather than strings?

fields.py:

from enum import Enum

from django.core.exceptions import ValidationError
from django.db import models


class EnumField(models.CharField):
    description = 'Enum with strictly typed choices'

    def __init__(self, enum_class, *args, **kwargs):
        self._enum_class = enum_class
        choices = []
        for enum in self._enum_class:
            title_case = enum.name.replace('_', ' ').title()
            entry = (enum, title_case)
            choices.append(entry)
        kwargs['choices'] = choices
        kwargs['blank'] = False  # blank doesn't make sense for enum's
        super().__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super().deconstruct()
        args.insert(0, self._enum_class)
        del kwargs['choices']
        return name, path, args, kwargs

    def from_db_values(self, value, expression, connection, context):
        return self.to_python(value)

    def to_python(self, value):
        if value is None or isinstance(value, self._enum_class):
            return value
        else:
            return self._parse_enum(value)

    def _parse_enum(self, value):
        try:
            enum = self._enum_class[value]
        except KeyError:
            raise ValidationError("Invalid type '{}' for {}".format(
                value, self._enum_class))
        else:
            return enum

    def get_prep_value(self, value):
        if value is None:
            return None
        elif isinstance(value, Enum):
            return value.name
        else:
            msg = "'{}' must have type {}".format(
                value, self._enum_class.__name__)
            if self.null:
                msg += ', or `None`'
            raise TypeError(msg)

    def get_choices(self, **kwargs):
        kwargs['include_blank'] = False  # Blank is not a valid option
        choices = super().get_choices(**kwargs)
        return choices
1

There are 1 answers

0
lgunsch On BEST ANSWER

After a lot of digging, I was able to answer my own question:

SubfieldBase has been deprecated, and will be removed in Django 1.10; which is why I left it out of the implementation above. However, it seems that what it does is still important. Adding the following method to replaces the functionality that SubfieldBase would have added.

def contribute_to_class(self, cls, name, **kwargs):
    super(EnumField, self).contribute_to_class(cls, name, **kwargs)
    setattr(cls, self.name, Creator(self))

The Creator descriptor is what calls to_python on attributes. If this didn't happen, querys on models would result in the EnumField fields in the model instances being simply strings, instead of Enum instances like I wanted.