fix _get_type - check for NoneType

piccolo_api/fastapi/endpoints.py

@@ -21,6 +21,15 @@
 ANNOTATIONS: t.DefaultDict = defaultdict(dict)
 
 
+try:
+    # Python 3.10 and above
+    from types import UnionType  # type: ignore
+except ImportError:
+
+    class UnionType:  # type: ignore
+        ...
+
+
 class HTTPMethod(str, Enum):
     get = "GET"
     delete = "DELETE"
@@ -83,25 +92,30 @@
 
     For example::
 
-        >>> get_type(Optional[int])
+        >>> _get_type(Optional[int])
         int
 
-        >>> get_type(int)
+        >>> _get_type(int | None)
         int
 
-        >>> get_type(list[str])
+        >>> _get_type(int)
+        int
+
+        >>> _get_type(list[str])
         list[str]
 
     """
     origin = t.get_origin(type_)
 
     # Note: even if `t.Optional` is passed in, the origin is still a
-    # `t.Union`.
-    if origin is t.Union:
-        args = t.get_args(type_)
+    # `t.Union` or `UnionType` depending on the Python version.
+    if any(origin is i for i in (t.Union, UnionType)):
+        union_args = t.get_args(type_)
+
+        NoneType = type(None)
 
-        if len(args) == 2 and None in args:
-            return [i for i in args if i is not None][0]
+        if len(union_args) == 2 and NoneType in union_args:
+            return [i for i in union_args if i is not NoneType][0]
 
     return type_
 

tests/fastapi/test_fastapi_endpoints.py

@@ -1,13 +1,16 @@
+import sys
+import typing as t
 from unittest import TestCase
 
+import pytest
 from fastapi import FastAPI
 from piccolo.columns import ForeignKey, Integer, Varchar
 from piccolo.columns.readable import Readable
 from piccolo.table import Table
 from starlette.testclient import TestClient
 
 from piccolo_api.crud.endpoints import PiccoloCRUD
-from piccolo_api.fastapi.endpoints import FastAPIWrapper
+from piccolo_api.fastapi.endpoints import FastAPIWrapper, _get_type
 
 
 class Movie(Table):
@@ -246,3 +249,27 @@ def test_patch(self):
         self.assertEqual(
             response.json(), {"id": 1, "name": "Star Wars", "rating": 90}
         )
+
+
+class TestGetType(TestCase):
+    def test_get_type(self):
+        """
+        If we pass in an optional type, it should return the non-optional type.
+        """
+        # Should return the underlying type, as they're all optional:
+        self.assertIs(_get_type(t.Optional[str]), str)
+        self.assertIs(_get_type(t.Optional[t.List[str]]), t.List[str])
+        self.assertIs(_get_type(t.Union[str, None]), str)
+
+        # Should be returned as is, because it's not optional:
+        self.assertIs(_get_type(t.List[str]), t.List[str])
+
+    @pytest.mark.skipif(
+        sys.version_info < (3, 10), reason="Union syntax not available"
+    )
+    def test_new_union_syntax(self):
+        """
+        Make sure it works with the new syntax added in Python 3.10.
+        """
+        self.assertIs(_get_type(str | None), str)  # type: ignore
+        self.assertIs(_get_type(None | str), str)  # type: ignore