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
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 thatSubfieldBase
would have added.The
Creator
descriptor is what callsto_python
on attributes. If this didn't happen, querys on models would result in theEnumField
fields in the model instances being simply strings, instead of Enum instances like I wanted.