From d7a81b185f1771b7c7945b202705dbfaeff011f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89tienne=20Noss?= <etienne.noss@wifirst.fr>
Date: Thu, 1 Aug 2019 10:45:44 +0200
Subject: [PATCH] fix Optional typing hint for SerializerMethodField

---
 src/drf_yasg/inspectors/field.py |  2 ++
 tests/test_schema_generator.py   | 41 ++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/src/drf_yasg/inspectors/field.py b/src/drf_yasg/inspectors/field.py
index c18de720..fa807184 100644
--- a/src/drf_yasg/inspectors/field.py
+++ b/src/drf_yasg/inspectors/field.py
@@ -610,6 +610,8 @@ def field_to_swagger_object(self, field, swagger_object_type, use_references, **
             # look for Python 3.5+ style type hinting of the return value
             hint_class = inspect_signature(method).return_annotation
 
+            if not inspect.isclass(hint_class) and hasattr(hint_class, '__args__'):
+                hint_class = hint_class.__args__[0]
             if inspect.isclass(hint_class) and not issubclass(hint_class, inspect._empty):
                 type_info = get_basic_type_info_from_hint(hint_class)
 
diff --git a/tests/test_schema_generator.py b/tests/test_schema_generator.py
index 24002384..b232db1b 100644
--- a/tests/test_schema_generator.py
+++ b/tests/test_schema_generator.py
@@ -1,4 +1,5 @@
 import json
+import sys
 from collections import OrderedDict
 
 import pytest
@@ -17,6 +18,11 @@
 from drf_yasg.generators import OpenAPISchemaGenerator
 from drf_yasg.utils import swagger_auto_schema
 
+try:
+    import typing
+except ImportError:
+    typing = None
+
 
 def test_schema_is_valid(swagger, codec_yaml):
     codec_yaml.encode(swagger)
@@ -293,3 +299,38 @@ class JSONViewSet(viewsets.ModelViewSet):
     swagger = generator.get_schema(None, True)
     property_schema = swagger["definitions"]["TestJSONField"]["properties"]["json"]
     assert property_schema == openapi.Schema(title='Json', type=openapi.TYPE_OBJECT)
+
+
+@pytest.mark.parametrize('py_type, expected_type', [
+    (str, openapi.TYPE_STRING),
+    (int, openapi.TYPE_INTEGER),
+    (float, openapi.TYPE_NUMBER),
+    (bool, openapi.TYPE_BOOLEAN),
+])
+@pytest.mark.skipif(typing is None or sys.version_info.major < 3, reason="typing not supported")
+def test_optional_return_type(py_type, expected_type):
+
+    class OptionalMethodSerializer(serializers.Serializer):
+        x = serializers.SerializerMethodField()
+
+        def get_x(self, instance):
+            pass
+
+        # Add the type annotation here in order to avoid a SyntaxError in py27
+        get_x.__annotations__["return"] = typing.Optional[py_type]
+
+    class OptionalMethodViewSet(viewsets.ViewSet):
+        @swagger_auto_schema(responses={200: openapi.Response("OK", OptionalMethodSerializer)})
+        def retrieve(self, request, pk=None):
+            return Response({'optional': None})
+
+    router = routers.DefaultRouter()
+    router.register(r'optional', OptionalMethodViewSet, **_basename_or_base_name('optional'))
+
+    generator = OpenAPISchemaGenerator(
+        info=openapi.Info(title='Test optional parameter', default_version='v1'),
+        patterns=router.urls
+    )
+    swagger = generator.get_schema(None, True)
+    property_schema = swagger["definitions"]["OptionalMethod"]["properties"]["x"]
+    assert property_schema == openapi.Schema(title='X', type=expected_type, readOnly=True)