Type hint support for Enums / TextChoices

[jalaziz]

Describe the bug

When using type hints with a serializer method field that returns a TextChoices/Enum type, we get the "unable to resolve type hint for function" warning.

While Enum value types are not guaranteed to be strings, Django provides the TextChoices and IntegerChoices base classes that have str and int values respectively.

It may be possible to support an even wider arrange of Enum by checking for both subclasses of both Enum and a basic supported type. This would cover enums defines as Foo(int, Enum), which is effectively how IntegerChoices is defined.

To Reproduce

class Bar(TextChoices):
    A = "a"
    B = "b"

class FooSerializer(Serializer):
    active_bars = serializers.SerializerMethodField()

    class Meta:
        fields = ('active_bars',)

    def get_active_bars(self, obj) -> list[Bar]:
        return [Bar.A]

Expected behavior

drf-spectacular should be correct identify the OpenAPI types for TextChoices (string) and IntegerChoices (number) enums.

[tfranzel]

Hi, i understand the goal you are trying to achieve. We have a whole heap of supported typings, but i think this approach is not correct.

• list[Bar] returns a list of Bar instances and not ['a', 'b'].
• In theory, you would need to do typing.Literal[*Bar.values] . We do support typing.Literal, but python does not allow * unpacking in types.

What does work is the following. Not looking that pretty, but semantically and syntactically correct:

@extend_schema_field(serializers.ListField(child=serializers.ChoiceField(choices=Bar.choices)))
def get_active_bars(self, obj):
        return [Bar.A]

EDIT: forgot to wrap the suggestion in a list

[jalaziz]

• list[Bar] returns a list of Bar instances and not ['a', 'b'].

Entirely true and fair.

• In theory, you would need to do typing.Literal[*Bar.values] . We do support typing.Literal, but python does not allow * unpacking in types.

Besides not actually being possible, it's also not semantically correct from an actual return type perspective if I'm not mistaken.

What does work is the following. Not looking that pretty, but semantically and syntactically correct:

@extend_schema_field(serializers.ListField(child=serializers.ChoiceField(choices=Bar.choices)))
def get_active_bars(self, obj):
        return [Bar.A]

Indeed that's both semantically and syntactically correct. However, since it's possible to represent as a serializer fields, wouldn't it be possible to support it from the type hinting?

I do realize it's difficult for any enum, but it seems like enums defined from Django's Choices might be safe to support?

The workaround is fine, just not pretty, and it just seems like something that would be nice out of the box. Happy to give a PR a shot myself if it's something you agree should be supported.

[tfranzel]

Turns out the truth is somewhere in between. I had a reason to not include it back then, but my understanding was incomplete. The implementation has had two goals:

• mypy should not complain about the hint
• DRF should not fail

class BarA(Enum):
    A = "a"
    B = "b"

# DRF serialization fails
def get_foo(self, obj) -> BarA:
        return BarA.A

# MYPY complains: BarA != str
def get_foo(self, obj) -> BarA:
        return BarA.A.value

# ---------------------------

class BarB(Choices):
    A = "a"
    B = "b"

# DRF serialization fails
def get_foo(self, obj) -> BarB:
        return BarB.A

# MYPY complains: BarB != str
def get_foo(self, obj) -> BarB:
        return BarB.A.value

# ---------------------------

class BarC(TextChoices):
    A = "a"
    B = "b"

# WORKS!
def get_foo(self, obj) -> BarC:
        return BarC.A

So TextChoices and IntegerChoices can work. Unfortunately Choices does not. Probably because of the MRO missing a basic type base class (class TextChoices(str, Choices)).

I'm unsure whether we should also support Enum and require the user to use # type: ignore if he uses mypy. Not sure if my OCD is talking here. What do you think?

[jalaziz]

From the Django docs, it really seems like Choices should always be used with a concrete value type.

Looking at the code for Choices, it doesn't seem like there's any enforcement of that though.

My personal opinion is:

1. TextChoices, IntegerChoices, and subclasses of (supported_base_type, Choices) should be supported.
2. If possible, subclasses of (supported_base_type, Enum) should be supported. This is, of course, if mypy and DRF serialization works.

My OCD is pushing me to encourage support for the most universal options, but realistically, I would imagine basic support for even just TextChoices and IntegerChoices would go a long way.

[jalaziz]

@tfranzel I checked out your branch to test out a few things, but seems like everything is passing. Was it fixed or is there a way for me to test out the scenarios you have outlined?

Happy to setup scenarios myself, but if you happen to have a "playground" setup, that would be awesome.

[tfranzel]

Looking at the code for Choices, it doesn't seem like there's any enforcement of that though.

that caught me too. So TextChoices, IntegerChoices are pretty much the only sensible "choices" for types.

yes the branch does cover everything including Enum. I wrote it and then discovered again, why i did not include it in the first place. i still need to reduce down the PR again to the TextChoices, IntegerChoices. This is the question i was posing about supporting somethings that does not work in practice.

[jalaziz]

that caught me too. So TextChoices, IntegerChoices are pretty much the only sensible "choices" for types.

The Django docs does show an example that uses date as the type. Not sure if that's really useful, but just thought I should mention it.

yes the branch does cover everything including Enum. I wrote it and then discovered again, why i did not include it in the first place. i still need to reduce down the PR again to the TextChoices, IntegerChoices. This is the question i was posing about supporting somethings that does not work in practice.

Makes sense, I mainly wanted to test if supporting (base_type, Enum) subclasses would be feasible, but wasn't sure how to easily reproduce the test scenarios you presented.

[tfranzel]

i refactored the change and it is now on master. please have a look. it should work for (base_type, Choices) where base_type is a limited set of basic python types (str, date, float, uuid, etc). of course that includes TextChoices.

[jalaziz]

@tfranzel tested using the master branch and indeed I no longer receive a system check error and the generated schema looks correct.

Thank you!!!

[ZackPlauche]

Guys, it seems like this is closed, but how do you actually type hint with a TextChoices subclass?

[tfranzel]

if you look at the merged PR you can see the added test using this feature:

drf-spectacular/tests/test_postprocessing.py

Lines 294 to 303 in 82c00f8

	def test_textchoice_annotation(no_warnings):
	class QualityChoices(TextChoices):
	GOOD = 'GOOD'
	BAD = 'BAD'
	class XSerializer(serializers.Serializer):
	quality_levels = serializers.SerializerMethodField()
	def get_quality_levels(self, obj) -> typing.List[QualityChoices]:
	return [QualityChoices.GOOD, QualityChoices.BAD] # pragma: no cover

[ArtyomShalagin]

Hey, as far as I can tell this works with list but not with dict. Like:

class Bar(TextChoices)
   A = "a"
   B = "b"

class FooSerializer(serializers.Serializer): 
   bar_list = serializers.SerializerMethodField()
   bar_dict = serializers.SerializerMethodField()
  
   def get_bar_list(self, obj) -> list[Bar]: 
      return [Bar.A]

   def get_bar_dict(self, obj) -> dict[Bar, int]: 
      return {Bar.A: 1} 

Generated docs look like this:

"bar_list": ["a"],
"bar_dict": {
   "additionalProp1": 0,
   "additionalProp2": 0,
   "additionalProp3": 0
},

The schema has Enum for list, but for dict it's just < * >: number($double). There is no "unable to resolve type hint" warning though. Is this expected behavior? That's on version 0.27.0