From bc4d52558bbf3d8a1311c23194123ea7517b2697 Mon Sep 17 00:00:00 2001
From: Dhaval Mehta <dhaval.mehta197@gmail.com>
Date: Wed, 29 Jan 2020 23:45:56 +0530
Subject: [PATCH] Schemas: Add mapping of type for ChoiceField. (#7161)

---
 rest_framework/schemas/openapi.py | 38 ++++++++++++++++++++++++++-----
 tests/schemas/test_openapi.py     | 24 ++++++++++++++++++-
 2 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/rest_framework/schemas/openapi.py b/rest_framework/schemas/openapi.py
index 6f3e854b0..fa887d63a 100644
--- a/rest_framework/schemas/openapi.py
+++ b/rest_framework/schemas/openapi.py
@@ -1,4 +1,6 @@
 import warnings
+from collections import OrderedDict
+from decimal import Decimal
 from operator import attrgetter
 from urllib.parse import urljoin
 
@@ -209,6 +211,34 @@ class AutoSchema(ViewInspector):
 
         return paginator.get_schema_operation_parameters(view)
 
+    def _map_choicefield(self, field):
+        choices = list(OrderedDict.fromkeys(field.choices))  # preserve order and remove duplicates
+        if all(isinstance(choice, bool) for choice in choices):
+            type = 'boolean'
+        elif all(isinstance(choice, int) for choice in choices):
+            type = 'integer'
+        elif all(isinstance(choice, (int, float, Decimal)) for choice in choices):  # `number` includes `integer`
+            # Ref: https://tools.ietf.org/html/draft-wright-json-schema-validation-00#section-5.21
+            type = 'number'
+        elif all(isinstance(choice, str) for choice in choices):
+            type = 'string'
+        else:
+            type = None
+
+        mapping = {
+            # The value of `enum` keyword MUST be an array and SHOULD be unique.
+            # Ref: https://tools.ietf.org/html/draft-wright-json-schema-validation-00#section-5.20
+            'enum': choices
+        }
+
+        # If We figured out `type` then and only then we should set it. It must be a string.
+        # Ref: https://swagger.io/docs/specification/data-models/data-types/#mixed-type
+        # It is optional but it can not be null.
+        # Ref: https://tools.ietf.org/html/draft-wright-json-schema-validation-00#section-5.21
+        if type:
+            mapping['type'] = type
+        return mapping
+
     def _map_field(self, field):
 
         # Nested Serializers, `many` or not.
@@ -242,15 +272,11 @@ class AutoSchema(ViewInspector):
         if isinstance(field, serializers.MultipleChoiceField):
             return {
                 'type': 'array',
-                'items': {
-                    'enum': list(field.choices)
-                },
+                'items': self._map_choicefield(field)
             }
 
         if isinstance(field, serializers.ChoiceField):
-            return {
-                'enum': list(field.choices),
-            }
+            return self._map_choicefield(field)
 
         # ListField.
         if isinstance(field, serializers.ListField):
diff --git a/tests/schemas/test_openapi.py b/tests/schemas/test_openapi.py
index dd3e87c7e..b4cb2823f 100644
--- a/tests/schemas/test_openapi.py
+++ b/tests/schemas/test_openapi.py
@@ -1,3 +1,5 @@
+import uuid
+
 import pytest
 from django.conf.urls import url
 from django.test import RequestFactory, TestCase, override_settings
@@ -44,6 +46,8 @@ class TestBasics(TestCase):
 
 class TestFieldMapping(TestCase):
     def test_list_field_mapping(self):
+        uuid1 = uuid.uuid4()
+        uuid2 = uuid.uuid4()
         inspector = AutoSchema()
         cases = [
             (serializers.ListField(), {'items': {}, 'type': 'array'}),
@@ -53,7 +57,25 @@ class TestFieldMapping(TestCase):
             (serializers.ListField(child=serializers.IntegerField(max_value=4294967295)),
              {'items': {'type': 'integer', 'maximum': 4294967295, 'format': 'int64'}, 'type': 'array'}),
             (serializers.ListField(child=serializers.ChoiceField(choices=[('a', 'Choice A'), ('b', 'Choice B')])),
-             {'items': {'enum': ['a', 'b']}, 'type': 'array'}),
+             {'items': {'enum': ['a', 'b'], 'type': 'string'}, 'type': 'array'}),
+            (serializers.ListField(child=serializers.ChoiceField(choices=[(1, 'One'), (2, 'Two')])),
+             {'items': {'enum': [1, 2], 'type': 'integer'}, 'type': 'array'}),
+            (serializers.ListField(child=serializers.ChoiceField(choices=[(1.1, 'First'), (2.2, 'Second')])),
+             {'items': {'enum': [1.1, 2.2], 'type': 'number'}, 'type': 'array'}),
+            (serializers.ListField(child=serializers.ChoiceField(choices=[(True, 'true'), (False, 'false')])),
+             {'items': {'enum': [True, False], 'type': 'boolean'}, 'type': 'array'}),
+            (serializers.ListField(child=serializers.ChoiceField(choices=[(uuid1, 'uuid1'), (uuid2, 'uuid2')])),
+             {'items': {'enum': [uuid1, uuid2]}, 'type': 'array'}),
+            (serializers.ListField(child=serializers.ChoiceField(choices=[(1, 'One'), ('a', 'Choice A')])),
+             {'items': {'enum': [1, 'a']}, 'type': 'array'}),
+            (serializers.ListField(child=serializers.ChoiceField(choices=[
+                (1, 'One'), ('a', 'Choice A'), (1.1, 'First'), (1.1, 'First'), (1, 'One'), ('a', 'Choice A'), (1, 'One')
+            ])),
+                {'items': {'enum': [1, 'a', 1.1]}, 'type': 'array'}),
+            (serializers.ListField(child=serializers.ChoiceField(choices=[
+                (1, 'One'), (2, 'Two'), (3, 'Three'), (2, 'Two'), (3, 'Three'), (1, 'One'),
+            ])),
+                {'items': {'enum': [1, 2, 3], 'type': 'integer'}, 'type': 'array'}),
             (serializers.IntegerField(min_value=2147483648),
              {'type': 'integer', 'minimum': 2147483648, 'format': 'int64'}),
         ]