From 7d3748fda9ddcba0acf005bc2d5f05b8321436a9 Mon Sep 17 00:00:00 2001
From: Marti Raudsepp <marti@juffo.org>
Date: Wed, 20 Jul 2022 18:50:46 +0300
Subject: [PATCH] Improvements

---
 src/drf_yasg/inspectors/field.py            | 19 ++++++++++++-------
 tests/test_get_basic_type_info_from_hint.py |  4 ++--
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/drf_yasg/inspectors/field.py b/src/drf_yasg/inspectors/field.py
index 65d399ee..b360ca76 100644
--- a/src/drf_yasg/inspectors/field.py
+++ b/src/drf_yasg/inspectors/field.py
@@ -2,6 +2,7 @@
 import inspect
 import logging
 import operator
+import types
 import typing
 import uuid
 from collections import OrderedDict
@@ -459,7 +460,11 @@ def decimal_return_type():
 
 
 def get_origin_type(hint_class):
-    return getattr(hint_class, '__origin__', None) or hint_class
+    # Added in Python 3.8
+    if hasattr(typing, 'get_origin'):
+        return typing.get_origin(hint_class)
+    else:
+        return getattr(hint_class, '__origin__', None) or hint_class
 
 
 def hint_class_issubclass(hint_class, check_class):
@@ -502,16 +507,16 @@ def inspect_collection_hint_class(hint_class):
 
 hinting_type_info.append(((typing.Sequence, typing.AbstractSet), inspect_collection_hint_class))
 
-# typing.UnionType was added in Python 3.10 for new PEP 604 pipe union syntax
-try:
-    from types import UnionType
-except ImportError:
-    UnionType = None
+# types.UnionType was added in Python 3.10 for new PEP 604 pipe union syntax
+if hasattr(types, 'UnionType'):
+    UNION_TYPES = (typing.Union, types.UnionType)
+else:
+    UNION_TYPES = (typing.Union,)
 
 
 def _get_union_types(hint_class):
     origin_type = get_origin_type(hint_class)
-    if origin_type is typing.Union or (UnionType is not None and origin_type is UnionType):
+    if origin_type in UNION_TYPES:
         return hint_class.__args__
 
 
diff --git a/tests/test_get_basic_type_info_from_hint.py b/tests/test_get_basic_type_info_from_hint.py
index eac39e93..85c9cccc 100644
--- a/tests/test_get_basic_type_info_from_hint.py
+++ b/tests/test_get_basic_type_info_from_hint.py
@@ -17,13 +17,13 @@
 
 python310_union_tests = []
 if sys.version_info >= (3, 10):
-    # # New PEP 604 union syntax in Python 3.10+
+    # New PEP 604 union syntax in Python 3.10+
     python310_union_tests = [
         (bool | None, {'type': openapi.TYPE_BOOLEAN, 'format': None, 'x-nullable': True}),
         (list[int] | None, {
             'type': openapi.TYPE_ARRAY, 'items': openapi.Items(openapi.TYPE_INTEGER), 'x-nullable': True
         }),
-        # Following cases are not 100% correct, but it should work somehow and not crash.
+        # Following case is not handled, but it should not crash.
         (int | float, None),
     ]