Last active
November 8, 2019 19:50
-
-
Save cdcarter/283ded66c7476cdc8bfc39f1e61dc9bd to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from enum import Enum | |
from collections import namedtuple | |
from django.utils.encoding import force_text | |
from django.db import models | |
from functools import partialmethod | |
class NamedTupleEnum(Enum): | |
__attrs__: namedtuple = None # Override this with your namedtuple factory. | |
def __new__(cls, value, attrs): | |
# We have to override __new__ so that only the first parameter is the enum member "value" | |
obj = object.__new__(cls) | |
obj._value_ = value | |
# See note in tests.py#StateEnum.RECONSTRUCT, but I'm leaving this out for now. | |
# if hasattr(attrs, "_replace"): # if its already a namedtuple instance | |
# obj._attrs = attrs # don't worry about constructing a new one. | |
obj._attrs = cls.__attrs__(*attrs) # we'll shove the NamedTuple into an attribute for later. | |
return obj | |
def __getattr__(self, name): | |
# TODO: should we actually put these attributes onto the Enum instance? | |
return getattr(self._attrs, name) | |
class BuildSource(NamedTupleEnum): | |
__attrs__ = namedtuple("BuildAttrs", "auto enqueue") | |
WEBHOOK = _("webhook"), (True, True) | |
MANUAL = _("manual"), (False, True) | |
CLI = _("cli"), (False, False) | |
LEGACY = _("legacy"), (False, True) | |
BuildSource.WEBHOOK.auto # => True |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# further usages of this with Django, less fun for y'all | |
class IntEnumField(models.IntegerField): | |
_enum: NamedTupleEnum | |
def __init__(self, *args, **kwargs): | |
self._enum = kwargs.pop("enum") | |
kwargs['choices'] = self._enum.choices | |
super().__init__(*args, **kwargs) | |
def deconstruct(self): | |
name, path, args, kwargs = super().deconstruct() | |
kwargs["enum"] = self._enum | |
del kwargs["choices"] | |
return name, path, args, kwargs | |
def from_db_value(self, value, *_): | |
if value is not None: | |
return self._enum(value) | |
return value | |
def get_prep_value(self, value): | |
if value is None: | |
return value | |
if isinstance(value, Enum): | |
return value.value | |
return int(value) | |
def _get_FIELD_display(self, cls): | |
value = getattr(cls, self.attname) | |
return force_text(value.label, strings_only=True) | |
def contribute_to_class(self, cls, name, private_only=False): | |
super().contribute_to_class(cls, name, private_only) | |
setattr(cls, "get_%s_display" % self.name, partialmethod(self._get_FIELD_display)) | |
def get_transform(self, lookup_name): | |
nt = self._enum.__attrs__ | |
if lookup_name in nt._fields: | |
return bool_enum_transform(lookup_name, self._enum) | |
return super().get_transform(lookup_name) | |
# noinspection PyPep8Naming | |
class classproperty: | |
""" lovingly lifted from Django HEAD. """ | |
def __init__(self, method=None): | |
self.fget = method | |
def __get__(self, instance, cls=None): | |
return self.fget(cls) | |
def getter(self, method): | |
self.fget = method | |
return self | |
class EnumChoices: | |
""" Mixin for Enums that are going to be used as `choices` for a Django field. """ | |
@classproperty | |
def choices(cls): | |
""" Generate a sequence of [db_value,label] pairs from the enum members, to be | |
used as the `choices` for a Django field. | |
The db_value will always be the enum member's value. The label will either | |
be the "label" property of the enum member, or the constant name (humanized). """ | |
# TODO: this collecting is not particularly pythonic, this should probably become a generator. | |
# TODO: enum needs a label property always for the display to work at this point. | |
ret_val = [] | |
for member in cls: | |
ret_member = [member.value] | |
if hasattr(member, "label"): | |
ret_member += [member.label] | |
else: | |
ret_member += [member.name.replace('_', ' ').title()] | |
ret_val += [ret_member] | |
return ret_val | |
class NamedTupleChoices(EnumChoices, NamedTupleEnum): | |
pass |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment