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 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.