From f05889292abf04dbf4148e575d35a56850d5e32a Mon Sep 17 00:00:00 2001
From: Nicolas Delaby <ticosax@free.fr>
Date: Fri, 22 Dec 2017 21:00:13 +0100
Subject: [PATCH] Add type inspection for PrimaryKeyRelatedField (#26)

Guess the type for PrimaryKeyRelatedField from the related ModelField.
---
 src/drf_yasg/generators.py |  7 +++----
 src/drf_yasg/utils.py      | 11 ++++++++++-
 tests/conftest.py          |  2 +-
 tests/reference.yaml       |  2 +-
 tox.ini                    |  2 +-
 5 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/drf_yasg/generators.py b/src/drf_yasg/generators.py
index 522d27d..6eeef64 100644
--- a/src/drf_yasg/generators.py
+++ b/src/drf_yasg/generators.py
@@ -1,7 +1,6 @@
 import re
 from collections import defaultdict, OrderedDict
 
-import django.db.models
 import uritemplate
 from coreapi.compat import force_text
 from rest_framework.schemas.generators import SchemaGenerator, EndpointEnumerator as _EndpointEnumerator
@@ -10,6 +9,7 @@ from rest_framework.schemas.inspectors import get_pk_description
 from . import openapi
 from .inspectors import SwaggerAutoSchema
 from .openapi import ReferenceResolver
+from .utils import get_schema_type_from_model_field
 
 PATH_PARAMETER_RE = re.compile(r'{(?P<parameter>\w+)}')
 
@@ -221,6 +221,8 @@ class OpenAPISchemaGenerator(object):
                     model_field = model._meta.get_field(variable)
                 except Exception:  # pragma: no cover
                     model_field = None
+                else:
+                    type = get_schema_type_from_model_field(model_field)
 
                 if model_field is not None and model_field.help_text:
                     description = force_text(model_field.help_text)
@@ -229,9 +231,6 @@ class OpenAPISchemaGenerator(object):
 
                 if hasattr(view_cls, 'lookup_value_regex') and getattr(view_cls, 'lookup_field', None) == variable:
                     pattern = view_cls.lookup_value_regex
-                elif isinstance(model_field, django.db.models.AutoField):
-                    type = openapi.TYPE_INTEGER
-
             field = openapi.Parameter(
                 name=variable,
                 required=True,
diff --git a/src/drf_yasg/utils.py b/src/drf_yasg/utils.py
index 3a74ae8..cab2beb 100644
--- a/src/drf_yasg/utils.py
+++ b/src/drf_yasg/utils.py
@@ -1,6 +1,7 @@
 from collections import OrderedDict
 
 from django.core.validators import RegexValidator
+from django.db import models
 from django.utils.encoding import force_text
 from rest_framework import serializers
 from rest_framework.mixins import RetrieveModelMixin, DestroyModelMixin, UpdateModelMixin
@@ -13,6 +14,12 @@ from .errors import SwaggerGenerationError
 no_body = object()
 
 
+def get_schema_type_from_model_field(model_field):
+    if isinstance(model_field, models.AutoField):
+        return openapi.TYPE_INTEGER
+    return openapi.TYPE_STRING
+
+
 def is_list_view(path, method, view):
     """Check if the given path/method appears to represent a list view (as opposed to a detail/instance view).
 
@@ -238,8 +245,10 @@ def serializer_field_to_swagger(field, swagger_object_type, definitions=None, **
             items=child_schema,
             unique_items=True,  # is this OK?
         )
+    elif isinstance(field, serializers.PrimaryKeyRelatedField):
+        model = field.queryset.model
+        return SwaggerType(type=get_schema_type_from_model_field(model._meta.pk))
     elif isinstance(field, serializers.RelatedField):
-        # TODO: infer type for PrimaryKeyRelatedField?
         return SwaggerType(type=openapi.TYPE_STRING)
     # ------ CHOICES
     elif isinstance(field, serializers.MultipleChoiceField):
diff --git a/tests/conftest.py b/tests/conftest.py
index d4efdb1..d390a8f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -40,7 +40,7 @@ def swagger_dict(generator):
 
 
 @pytest.fixture
-def validate_schema():
+def validate_schema(db):
     def validate_schema(swagger):
         from flex.core import parse as validate_flex
         from swagger_spec_validator.validator20 import validate_spec as validate_ssv
diff --git a/tests/reference.yaml b/tests/reference.yaml
index 24d4af2..9c04800 100644
--- a/tests/reference.yaml
+++ b/tests/reference.yaml
@@ -1042,7 +1042,7 @@ definitions:
       snippets:
         type: array
         items:
-          type: string
+          type: integer
         uniqueItems: true
       last_connected_ip:
         description: i'm out of ideas
diff --git a/tox.ini b/tox.ini
index a254d67..5d5c09f 100644
--- a/tox.ini
+++ b/tox.ini
@@ -30,7 +30,7 @@ deps =
     -rrequirements/test.txt
 
 commands =
-    pytest --cov-config .coveragerc --cov-append --cov
+    pytest --cov-config .coveragerc --cov-append --cov {posargs}
 
 [testenv:py36-drfmaster]
 pip_pre = True