From 27f32faee4dd96c5baee920d3a8f42ab4364782b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jeremy=20Lain=C3=A9?= <jeremy.laine@m4x.org>
Date: Thu, 1 Feb 2018 16:14:35 +0100
Subject: [PATCH] Fix schema generation for PrimaryKeyRelatedField (#5764)

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 413b9c0ac..8ef783683 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 1cbee0695..1abe5d2a1 100644
--- a/tests/test_schemas.py
+++ b/tests/test_schemas.py
@@ -23,7 +23,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()
 
@@ -556,6 +556,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):