How can I subclass Django TextChoices to add additional attributes?

Question

I want to use Django 3.0 TextChoices for a models.CharField choices option and also as an Enum in other places in my code.

This is a simplified version of my code:

from django.db import models

class ValueTypeOriginal(models.TextChoices):
    # I know Django will add the label for me, I am just being explicit
    BOOLEAN = 'boolean', 'Boolean'

class Template(models.Model):
    value_type = models.CharField(choices=ValueTypeOriginal.choices)

I am wanting to add an additional attribute to the enum members, something so that this call

>>> ValueType.BOOLEAN.native_type
bool

works.

bool here isn't a string, but the built-in python function .

This blog post described doing something like that with Enum by overriding __new__ .

class Direction(Enum):
    left = 37, (-1, 0)
    up = 38, (0, -1)
    right = 39, (1, 0)
    down = 40, (0, 1)

    def __new__(cls, keycode, vector):
        obj = object.__new__(cls)
        obj._value_ = keycode
        obj.vector = vector
        return obj

Based on that I tried:

class ValueTypeModified(models.TextChoices):
     BOOLEAN = ('boolean', bool), 'Boolean'
 
     def __new__(cls, value):
         obj = str.__new__(cls, value)
         obj._value_, obj.native_type = value
         return obj

That almost works. I get access to the unique TextChoices attributes like .choices , and I have the attribute .native_type but string comparison doesn't work like it should.

>>> ValueTypeOriginal.BOOLEAN == 'boolean'
True
>>> ValueTypeModified.BOOLEAN == 'boolean'
False

I think I am misunderstanding the __new__ method, but I'm stumped as to what I should be doing differently.

Update

In response to Ethan Furman's answer I tried

class ValueType(TextChoices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value, type):
        obj = str.__new__(value)
        obj._value_ = value
        obj.native_type = type
        return obj

but get

TypeError: __new__() missing 1 required positional argument: 'type'

So I went back to

class ValueType(TextChoices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value):
        obj = str.__new__(value)
        obj._value_ = value[0]
        obj.native_type = value[1]
        return obj

but I get:

TypeError: str.__new__(X): X is not a type object (tuple)

so then

class ValueType(TextChoices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value):
        obj = str.__new__(cls)
        obj._value_ = value[0]
        obj.native_type = value[1]
        return obj

gets me back where I started with direct string comparison failing

>>> ValueTypeOriginal.BOOLEAN == 'boolean'
True
>>> ValueType.BOOLEAN == 'boolean'
False

However,

>>> ValueType.BOOLEAN.value == 'boolean'
True

So the right value seems to get there, but the enum member itself isn't evaluating like a ValueType(str, Enum) but instead like ValueType(Enum) on comparisons.

Update #2

I've now tried:

class ValueType(TextChoices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value):
        obj = str.__new__(cls, value)
        obj._value_ = value[0]
        obj.native_type = value[1]
        return obj

class ValueType(str, Choices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value):
        obj = str.__new__(cls)
        obj._value_ = value[0]
        obj.native_type = value[1]
        return obj

and just to be safe

class ValueType(TextChoices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value):
        obj = super().__new__(cls, value)
        obj._value_ = value[0]
        obj.native_type = value[1]
        return obj

but none give me direct string comparison as expected.

Update #3 I finally understood what Ethan Furman was telling me to do.

Solution:

class ValueType(TextChoices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value):
        obj = str.__new__(cls, value[0])
        obj._value_ = value[0]
        obj.native_type = value[1]
        return obj

Answer 1

You are almost there. Part of the difficulty is only the first of the two arguments, ie the tuple ('boolean', bool) , is being passed to the Enum machinery.

So, we have two choices:

• keep the tuple as-is and use index-access (your current working solution):

  def __new__(cls, value): # value[0] is 'boolean'; value[1] is bool

• name the arguments in the __new__ header:

  def __new__(cls, svalue, type): # value is split into named arguments

Note that I changed the names slightly to hopefully help avoid confusion.

Putting it all together, your final method should look like (using the second option above):

def __new__(cls, svalue, type):
    obj = str.__new__(cls, svalue)
    obj._value_ = svalue
    obj.native_type = type
    return obj

Note :

The first argument to __new__ is the class of the instance you are trying to create -- typically the same class that the __new__ method is defined in. Even though it doesn't look like it, __new__ is a classmethod -- it's just special-cased to not require the classmethod decorator.

Answer 2

This is what I did using some metaclass magick:

class Mwahaha(type(models.TextChoices)):
    def __new__(metacls, classname, bases, classdict):
        native_types = {member: classdict[member][1] for member in classdict._member_names}
        classdict._member_names.clear()
        for member in native_types.keys():
            val = classdict[member][0]
            del classdict[member]
            classdict[member] = val
        cls = super().__new__(metacls, classname, bases, classdict)
        for member, n_t in native_types.items():
            getattr(o, member).native_type = n_t
        return cls

and have your class look like

class ValueTypeModified(models.TextChoices, metaclass=Mwahaha):
     BOOLEAN = ('boolean', 'Boolean'), bool

The clear and del are necessary to get around some of the Enum._dict protections preventing overriding enums or attributes. However, that's exactly what we want to do here.

There's probably an easier way to do this without resorting to the metaclass but I was primed and ready to go that route ¯\\_(ツ)_/¯

Answer 3

I'm posting my own answer so I can explain what I learned thanks to @Ethan Furman's awesome help!

From your code it looks like value is ('boolean', bool) , so when you do

obj = str.__new__(cls, value)

obj ends up being "('boolean', bool)"

That means this will work, even though it wasn't my intention

>>> ValueType.BOOLEAN == str(('boolean', bool))
True

Similarly, if I don't pass value at all to str.__new__ constructor (ie str.__new__(cls) ), then obj ends up being the empty string '' , just like calling str() with no arguments.

That means this would work, even though it wasn't my intention:

>>> ValueTypeEmptyString.BOOLEAN == ''
True

In the end, it really was my misunderstanding of the __new__ dunder method . Since I am doing a str.__new__ call not just a generic object.__new__ call, the first argument should be str itself or a subclass of str . In my case TextChoices is a subclass of str, so ValueType is also a subclass of str and can be the first argument to the str.__new__ method.

Then, as the docs for __new__ explain,

The remaining arguments are those passed to the object constructor expression (the call to the class).

Or in other words, I can think of the remaining arguments as feeding directly into a str() call. Since I don't want to stringify the whole tuple, but just the first element of that tuple, I should pass only the first element to the str.__new__ call.

So putting it all together:

class ValueType(TextChoices):
    BOOLEAN = (('boolean', bool), 'Boolean')
    
    def __new__(cls, value):
        # cls is <enum 'ValueType'>, a subtype of str, while value[0] is `'boolean'`
        obj = str.__new__(cls, value[0])
        # value[0] is again `'boolean'`
        obj._value_ = value[0]
        # value[1] is `bool`
        obj.native_type = value[1]
        return obj

The means by which the ChoicesMeta metaclass handles the addition of the passed Boolean label in the outer tuple, along with the other metaclass magic for .choices , isn't totally clear to me, but now I at least have the "working code" I was looking for and

>>> ValueType.BOOLEAN == 'boolean'
True

works.