From ee46f59fb1ee8b3ed9d0743851e7773c2d09e350 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cristi=20V=C3=AEjdea?= <cristi@cvjd.me>
Date: Mon, 26 Feb 2018 20:31:12 +0200
Subject: [PATCH] Infer ChoiceField type from model field when in
 ModelSerializer

Fixes part of issue #69
---
 .idea/drf-yasg.iml                            |  6 ++--
 docs/changelog.rst                            | 11 +++++++
 src/drf_yasg/inspectors/field.py              | 29 +++++++++++++------
 .../migrations/0002_article_article_type.py   | 18 ++++++++++++
 testproj/articles/models.py                   |  4 +++
 testproj/articles/serializers.py              |  2 +-
 tests/reference.yaml                          | 10 +++++++
 7 files changed, 67 insertions(+), 13 deletions(-)
 create mode 100644 testproj/articles/migrations/0002_article_article_type.py

diff --git a/.idea/drf-yasg.iml b/.idea/drf-yasg.iml
index c0ef6ba..308af4f 100644
--- a/.idea/drf-yasg.iml
+++ b/.idea/drf-yasg.iml
@@ -13,9 +13,9 @@
   </component>
   <component name="NewModuleRootManager">
     <content url="file://$MODULE_DIR$">
-      <sourceFolder url="file://$MODULE_DIR$" isTestSource="false" />
       <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
       <sourceFolder url="file://$MODULE_DIR$/testproj" isTestSource="false" />
+      <excludeFolder url="file://$MODULE_DIR$/venv" />
     </content>
     <orderEntry type="jdk" jdkName="Python 3.6 (drf-yasg)" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
@@ -24,7 +24,7 @@
     <option name="TEMPLATE_CONFIGURATION" value="Django" />
     <option name="TEMPLATE_FOLDERS">
       <list>
-        <option value="$MODULE_DIR$/drf_yasg/templates" />
+        <option value="$MODULE_DIR$/src/drf_yasg/templates" />
       </list>
     </option>
   </component>
@@ -32,4 +32,4 @@
     <option name="projectConfiguration" value="py.test" />
     <option name="PROJECT_TEST_RUNNER" value="py.test" />
   </component>
-</module>
+</module>
\ No newline at end of file
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 0c11374..a0fdf65 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -3,6 +3,17 @@ Changelog
 #########
 
 
+*********
+**1.4.4**
+*********
+
+*Release date: Feb 26, 2018*
+
+- **IMPROVED:** ``type`` for ``ChoiceField`` generated by a ``ModelSerializer`` from a model field with ``choices=...``
+  will now be set according to the associated model field (:issue:`69`)
+- **FIXED:** ``lookup_field`` and ``lookup_value_regex`` on the same ``ViewSet``  will no longer trigger an exception
+  (:issue:`68`)
+
 *********
 **1.4.3**
 *********
diff --git a/src/drf_yasg/inspectors/field.py b/src/drf_yasg/inspectors/field.py
index 4005d1a..68965ae 100644
--- a/src/drf_yasg/inspectors/field.py
+++ b/src/drf_yasg/inspectors/field.py
@@ -397,16 +397,27 @@ class ChoiceFieldInspector(FieldInspector):
     def field_to_swagger_object(self, field, swagger_object_type, use_references, **kwargs):
         SwaggerType, ChildSwaggerType = self._get_partial_types(field, swagger_object_type, use_references, **kwargs)
 
-        if isinstance(field, serializers.MultipleChoiceField):
-            return SwaggerType(
-                type=openapi.TYPE_ARRAY,
-                items=ChildSwaggerType(
-                    type=openapi.TYPE_STRING,
-                    enum=list(field.choices.keys())
+        if isinstance(field, serializers.ChoiceField):
+            enum_type = openapi.TYPE_STRING
+
+            # for ModelSerializer, try to infer the type from the associated model field
+            serializer = get_parent_serializer(field)
+            if isinstance(serializer, serializers.ModelSerializer):
+                model = getattr(getattr(serializer, 'Meta'), 'model')
+                model_field = get_model_field(model, field.source)
+                if model_field:
+                    enum_type = get_basic_type_info(model_field).get('type', enum_type)
+
+            if isinstance(field, serializers.MultipleChoiceField):
+                return SwaggerType(
+                    type=openapi.TYPE_ARRAY,
+                    items=ChildSwaggerType(
+                        type=enum_type,
+                        enum=list(field.choices.keys())
+                    )
                 )
-            )
-        elif isinstance(field, serializers.ChoiceField):
-            return SwaggerType(type=openapi.TYPE_STRING, enum=list(field.choices.keys()))
+
+            return SwaggerType(type=enum_type, enum=list(field.choices.keys()))
 
         return NotHandled
 
diff --git a/testproj/articles/migrations/0002_article_article_type.py b/testproj/articles/migrations/0002_article_article_type.py
new file mode 100644
index 0000000..abf6156
--- /dev/null
+++ b/testproj/articles/migrations/0002_article_article_type.py
@@ -0,0 +1,18 @@
+# Generated by Django 2.0.1 on 2018-02-26 18:32
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('articles', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='article',
+            name='article_type',
+            field=models.PositiveSmallIntegerField(choices=[(1, 'first'), (2, 'second'), (3, 'third'), (7, 'seven'), (8, 'eight')], help_text='IntegerField declared on model with choices=(...) and exposed via ModelSerializer', null=True),
+        ),
+    ]
diff --git a/testproj/articles/models.py b/testproj/articles/models.py
index 24878b8..ca6f453 100644
--- a/testproj/articles/models.py
+++ b/testproj/articles/models.py
@@ -8,5 +8,9 @@ class Article(models.Model):
     date_created = models.DateTimeField(auto_now_add=True)
     date_modified = models.DateTimeField(auto_now=True)
     author = models.ForeignKey('auth.User', related_name='articles', on_delete=models.CASCADE)
+    article_type = models.PositiveSmallIntegerField(
+        help_text="IntegerField declared on model with choices=(...) and exposed via ModelSerializer",
+        choices=((1, "first"), (2, "second"), (3, "third"), (7, "seven"), (8, "eight")), null=True
+    )
 
     cover = models.ImageField(upload_to='article/original/', blank=True)
diff --git a/testproj/articles/serializers.py b/testproj/articles/serializers.py
index 6a213f6..c447b71 100644
--- a/testproj/articles/serializers.py
+++ b/testproj/articles/serializers.py
@@ -16,7 +16,7 @@ class ArticleSerializer(serializers.ModelSerializer):
     class Meta:
         model = Article
         fields = ('title', 'author', 'body', 'slug', 'date_created', 'date_modified',
-                  'references', 'uuid', 'cover', 'cover_name')
+                  'references', 'uuid', 'cover', 'cover_name', 'article_type')
         read_only_fields = ('date_created', 'date_modified',
                             'references', 'uuid', 'cover_name')
         lookup_field = 'slug'
diff --git a/tests/reference.yaml b/tests/reference.yaml
index e82e99e..78cf596 100644
--- a/tests/reference.yaml
+++ b/tests/reference.yaml
@@ -591,6 +591,16 @@ definitions:
       cover_name:
         type: string
         readOnly: true
+      article_type:
+        description: IntegerField declared on model with choices=(...) and exposed
+          via ModelSerializer
+        type: integer
+        enum:
+          - 1
+          - 2
+          - 3
+          - 7
+          - 8
   Project:
     required:
       - projectName