From 5bf630ba2954e3cf0fc9b2b28c287c894a64afa1 Mon Sep 17 00:00:00 2001 From: Adam DePue Date: Wed, 3 Aug 2016 17:26:28 +0000 Subject: [PATCH] Add contrib methods to make setting work correctly --- src/django_richenum/models/fields.py | 29 ++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/src/django_richenum/models/fields.py b/src/django_richenum/models/fields.py index 084334c..6da93a1 100644 --- a/src/django_richenum/models/fields.py +++ b/src/django_richenum/models/fields.py @@ -5,12 +5,35 @@ from richenum import RichEnumValue, OrderedRichEnumValue +# https://github.com/django/django/blob/64200c14e0072ba0ffef86da46b2ea82fd1e019a/django/db/models/fields/subclassing.py#L31-L44 +class Creator(object): + """ + A placeholder class that provides a way to set the attribute on the model. + """ + def __init__(self, field): + self.field = field + + def __get__(self, obj, type=None): + if obj is None: + return self + return obj.__dict__[self.field.name] + + def __set__(self, obj, value): + obj.__dict__[self.field.name] = self.field.to_python(value) + + class IndexEnumField(models.IntegerField): '''Store ints in DB, but expose OrderedRichEnumValues in Python. ''' description = 'Efficient storage for OrderedRichEnums' + def contribute_to_class(self, cls, name, **kwargs): + super(IndexEnumField, self).contribute_to_class(cls, name, **kwargs) + + # Add Creator descriptor to allow the field to be set directly + setattr(cls, self.name, Creator(self)) + def __init__(self, enum, *args, **kwargs): if not hasattr(enum, 'from_index'): raise TypeError("%s doesn't support index-based lookup." % enum) @@ -110,6 +133,12 @@ class CanonicalNameEnumField(models.CharField): ''' description = 'Storage for RichEnums' + def contribute_to_class(self, cls, name, **kwargs): + super(CanonicalNameEnumField, self).contribute_to_class(cls, name, **kwargs) + + # Add Creator descriptor to allow the field to be set directly + setattr(cls, self.name, Creator(self)) + def __init__(self, enum, *args, **kwargs): if not hasattr(enum, 'from_canonical'): raise TypeError("%s doesn't support canonical_name lookup." % enum)