함수 스키마 변경: 유니온 및 리터럴 타입 처리 수정

Author: comfuture

function_schema/core.py

@@ -8,7 +8,6 @@
 py_310 = packaging.version.parse("3.10")
 
 if current_version >= py_310:
-    import types
     from types import UnionType
 else:
     UnionType = typing.Union  # type: ignore
@@ -154,13 +153,17 @@ def guess_type(
 
     origin = typing.get_origin(T)
 
+    if origin is typing.Annotated:
+        return guess_type(typing.get_args(T)[0])
+
     # hacking around typing modules, `typing.Union` and `types.UnitonType`
-    if origin is typing.Union or origin is UnionType:
+    if origin in [typing.Union, UnionType]:
         union_types = [t for t in typing.get_args(T) if t is not type(None)]
-        _types = []
-        for union_type in union_types:
-            _types.append(guess_type(union_type))
-        _types = [t for t in _types if t is not None]  # exclude None
+        _types = [
+            guess_type(union_type)
+            for union_type in union_types
+            if guess_type(union_type) is not None
+        ]
 
         # number contains integer in JSON schema
         if "number" in _types and "integer" in _types:
@@ -171,7 +174,8 @@ def guess_type(
         return _types
 
     if origin is typing.Literal:
-        return guess_type(typing.Union[tuple(type(arg) for arg in typing.get_args(T))])
+        type_args = typing.Union[tuple(type(arg) for arg in typing.get_args(T))]
+        return guess_type(type_args)
     elif origin is list or origin is tuple:
         return "array"
     elif origin is dict:

test/test_guess_type.py

@@ -98,17 +98,21 @@ def test_union_type():
 
 def test_literal_type():
     """Test literal type"""
-    assert guess_type(typing.Literal["a"]) == "string"
-    assert guess_type(typing.Literal[1]) == "integer"
+    assert guess_type(typing.Literal["a"]) == "string", "should be string"
+    assert guess_type(typing.Literal[1]) == "integer", "should be integer"
+    assert guess_type(typing.Literal[1.0]) == "number", "should be number"
 
-    assert set(guess_type(typing.Literal["a", 1, None])) == {"string", "integer"}
+    assert set(guess_type(typing.Literal["a", 1, None])) == {
+        "string",
+        "integer",
+    }, "should be string or integer, omit None"
 
     assert set(guess_type(typing.Literal["a", 1])) == {"string", "integer"}
-    assert set(guess_type(typing.Literal["a", 1.0])) == {"string", "integer"}
+    assert set(guess_type(typing.Literal["a", 1.0])) == {"string", "number"}
     assert set(guess_type(typing.Literal["a", 1.1])) == {"string", "number"}
     assert set(guess_type(typing.Literal["a", 1, 1.0])) == {
         "string",
         "number",
-    }  # XXX should be ["string", "integer", "number"] ?
+    }, "should omit integer if number is present"
 
     assert set(guess_type(typing.Literal["a", 1, 1.0, None])) == {"string", "number"}