add Enum support in type hints #492

drf_spectacular/plumbing.py

@@ -969,6 +969,11 @@ def resolve_type_hint(hint):
         if all(type(args[0]) is type(choice) for choice in args):
             schema.update(build_basic_type(type(args[0])))
         return schema
+    elif inspect.isclass(hint) and issubclass(hint, Choices):
+        return {
+            'enum': [item.value for item in hint],
+            **build_basic_type([t for t in hint.__mro__ if is_basic_type(t)][0])
+        }
     elif hasattr(typing, 'TypedDict') and isinstance(hint, typing._TypedDictMeta):
         return build_object_type(
             properties={

tests/test_plumbing.py

@@ -6,6 +6,7 @@
 from datetime import datetime
 
 import pytest
+from django import __version__ as DJANGO_VERSION
 from django.conf.urls import include
 from django.db import models
 from django.urls import re_path
@@ -135,6 +136,18 @@ class NamedTupleB(typing.NamedTuple):
     )
 ]
 
+if DJANGO_VERSION > '3':
+    from django.db.models.enums import TextChoices  # only available in Django>3
+
+    class LanguageChoices(TextChoices):
+        EN = 'en'
+        DE = 'de'
+
+    TYPE_HINT_TEST_PARAMS.append((
+        LanguageChoices,
+        {'enum': ['en', 'de'], 'type': 'string'}
+    ))
+
 if sys.version_info >= (3, 7):
     TYPE_HINT_TEST_PARAMS.append((
         typing.Iterable[collections.namedtuple("NamedTupleA", "a, b")],  # noqa

tests/test_postprocessing.py

@@ -1,15 +1,17 @@
+import typing
 from enum import Enum
 from unittest import mock
 
+import pytest
 from django import __version__ as DJANGO_VERSION
 from rest_framework import generics, mixins, serializers, viewsets
 from rest_framework.decorators import action
 from rest_framework.views import APIView
 
 try:
-    from django.db.models.enums import Choices
+    from django.db.models.enums import TextChoices
 except ImportError:
-    Choices = object  # type: ignore  # django < 3.0 handling
+    TextChoices = object  # type: ignore  # django < 3.0 handling
 
 from drf_spectacular.plumbing import list_hash, load_enum_name_overrides
 from drf_spectacular.utils import OpenApiParameter, extend_schema
@@ -35,7 +37,7 @@ class LanguageEnum(Enum):
     EN = 'en'
 
 
-class LanguageChoices(Choices):
+class LanguageChoices(TextChoices):
     EN = 'en'
 
 
@@ -190,3 +192,29 @@ def test_enum_override_variations(no_warnings):
 def test_enum_override_loading_fail(capsys):
     load_enum_name_overrides()
     assert 'unable to load choice override for LanguageEnum' in capsys.readouterr().err
+
+
+@pytest.mark.skipif(DJANGO_VERSION < '3', reason='Not available before Django 3.0')
+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
+
+    class XAPIView(APIView):
+        @extend_schema(responses=XSerializer)
+        def get(self, request):
+            pass  # pragma: no cover
+
+    schema = generate_schema('x', view=XAPIView)
+    assert 'QualityLevelsEnum' in schema['components']['schemas']
+    assert schema['components']['schemas']['X']['properties']['quality_levels'] == {
+        'type': 'array',
+        'items': {'$ref': '#/components/schemas/QualityLevelsEnum'},
+        'readOnly': True
+    }