From 6e4ad2565cdca01ccaed5fac65aadca2269b63f6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jeremy=20Lain=C3=A9?= <jeremy.laine@m4x.org>
Date: Tue, 23 Jan 2018 15:46:40 +0100
Subject: [PATCH] Fix schema generation for PrimaryKeyRelatedField

By default all subclasses of RelatedField are output as string fields in
the schema, which works well for StringRelatedField, SlugRelatedField or
HyperlinkedRelatedField.

Handle the common case of a PrimaryKeyRelatedField pointing to an
AutoField.
---
 rest_framework/schemas/inspectors.py |  8 +++++
 tests/test_schemas.py                | 47 +++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/rest_framework/schemas/inspectors.py b/rest_framework/schemas/inspectors.py
index 2afef7a04..796b095a6 100644
--- a/rest_framework/schemas/inspectors.py
+++ b/rest_framework/schemas/inspectors.py
@@ -50,6 +50,14 @@ def field_to_schema(field):
             title=title,
             description=description
         )
+    elif isinstance(field, serializers.PrimaryKeyRelatedField):
+        schema_cls = coreschema.String
+        model = getattr(field.queryset, 'model', None)
+        if model is not None:
+            model_field = model._meta.pk
+            if isinstance(model_field, models.AutoField):
+                schema_cls = coreschema.Integer
+        return schema_cls(title=title, description=description)
     elif isinstance(field, serializers.RelatedField):
         return coreschema.String(title=title, description=description)
     elif isinstance(field, serializers.MultipleChoiceField):
diff --git a/tests/test_schemas.py b/tests/test_schemas.py
index 34cb20798..7d580c85a 100644
--- a/tests/test_schemas.py
+++ b/tests/test_schemas.py
@@ -25,7 +25,7 @@ from rest_framework.utils import formatting
 from rest_framework.views import APIView
 from rest_framework.viewsets import GenericViewSet, ModelViewSet
 
-from .models import BasicModel
+from .models import BasicModel, ForeignKeySource
 
 factory = APIRequestFactory()
 
@@ -558,6 +558,51 @@ class TestSchemaGeneratorWithRestrictedViewSets(TestCase):
         assert schema == expected
 
 
+class ForeignKeySourceSerializer(serializers.ModelSerializer):
+    class Meta:
+        model = ForeignKeySource
+        fields = ('id', 'name', 'target')
+
+
+class ForeignKeySourceView(generics.CreateAPIView):
+    queryset = ForeignKeySource.objects.all()
+    serializer_class = ForeignKeySourceSerializer
+
+
+@unittest.skipUnless(coreapi, 'coreapi is not installed')
+class TestSchemaGeneratorWithForeignKey(TestCase):
+    def setUp(self):
+        self.patterns = [
+            url(r'^example/?$', ForeignKeySourceView.as_view()),
+        ]
+
+    def test_schema_for_regular_views(self):
+        """
+        Ensure that AutoField foreign keys are output as Integer.
+        """
+        generator = SchemaGenerator(title='Example API', patterns=self.patterns)
+        schema = generator.get_schema()
+
+        expected = coreapi.Document(
+            url='',
+            title='Example API',
+            content={
+                'example': {
+                    'create': coreapi.Link(
+                        url='/example/',
+                        action='post',
+                        encoding='application/json',
+                        fields=[
+                            coreapi.Field('name', required=True, location='form', schema=coreschema.String(title='Name')),
+                            coreapi.Field('target', required=True, location='form', schema=coreschema.Integer(description='Target', title='Target')),
+                        ]
+                    )
+                }
+            }
+        )
+        assert schema == expected
+
+
 @unittest.skipUnless(coreapi, 'coreapi is not installed')
 class Test4605Regression(TestCase):
     def test_4605_regression(self):