Merge pull request #4 from comfuture/literal-type

Fix enum handling in get_function_schema

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
@@ -114,7 +113,7 @@ def get_function_schema(
         }
 
         if enum_ is not None:
-            schema["properties"][name]["enum"] = [t for t in enum_]
+            schema["properties"][name]["enum"] = [t for t in enum_ if t is not None]
 
         if default_value is not inspect._empty:
             schema["properties"][name]["default"] = default_value
@@ -126,8 +125,14 @@ def get_function_schema(
         ):
             schema["required"].append(name)
 
+        if typing.get_origin(T) is typing.Literal:
+            if all(typing.get_args(T)):
+                schema["required"].append(name)
+
     parms_key = "input_schema" if format == "claude" else "parameters"
 
+    schema["required"] = list(set(schema["required"]))
+
     return {
         "name": func.__name__,
         "description": inspect.getdoc(func),
@@ -148,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:
@@ -165,7 +174,8 @@ def guess_type(
         return _types
 
     if origin is typing.Literal:
-        return "string"
+        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:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "function-schema"
-version = "0.3.3"
+version = "0.3.4"
 requires-python = ">= 3.9"
 description = "A small utility to generate JSON schemas for python functions."
 readme = "README.md"

test/test_guess_type.py

@@ -98,5 +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]) == "string"
+    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",
+    }, "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", "number"}
+    assert set(guess_type(typing.Literal["a", 1.1])) == {"string", "number"}
+    assert set(guess_type(typing.Literal["a", 1, 1.0])) == {
+        "string",
+        "number",
+    }, "should omit integer if number is present"
+
+    assert set(guess_type(typing.Literal["a", 1, 1.0, None])) == {"string", "number"}

test/test_schema.py

@@ -47,10 +47,10 @@ def func1(a: int, b: str, c: float = 1.0):
     assert (
         schema["parameters"]["properties"]["c"]["default"] == 1.0
     ), "c should have a default value of 1.0"
-    assert schema["parameters"]["required"] == [
-        "a",
-        "b",
-    ], "parameters with no default value should be required"
+    assert (
+        "a" in schema["parameters"]["required"]
+        and "b" in schema["parameters"]["required"]
+    ), "parameters with no default value should be required"
 
 
 def test_annotated_function():
@@ -78,10 +78,10 @@ def func1(
         schema["parameters"]["properties"]["b"]["description"] == "A string parameter"
     ), "parameter b should have a description"
 
-    assert schema["parameters"]["required"] == [
-        "a",
-        "b",
-    ], "parameters with no default value should be required"
+    assert (
+        "a" in schema["parameters"]["required"]
+        and "b" in schema["parameters"]["required"]
+    ), "parameters with no default value should be required"
 
 
 def test_annotated_function_with_enum():
@@ -116,7 +116,6 @@ def func1(animal: Annotated[Literal["Cat", "Dog"], "The animal you want to pet"]
         ...
 
     schema = get_function_schema(func1)
-    print(schema)
     assert (
         schema["parameters"]["properties"]["animal"]["type"] == "string"
     ), "parameter animal should be a string"
@@ -138,3 +137,37 @@ def func2(animal: Literal["Cat", "Dog"]):
         "Cat",
         "Dog",
     ], "parameter animal should have an enum"
+
+    assert "animal" in schema["parameters"]["required"], "animal is required"
+
+    def func3(animal: Literal["Cat", "Dog", 1]):
+        """My function"""
+        ...
+
+    schema = get_function_schema(func3)
+    assert schema["parameters"]["properties"]["animal"]["type"] == [
+        "string",
+        "integer",
+    ], "parameter animal should be a string"
+
+    def func4(animal: Literal["Cat", "Dog", 1, None]):
+        """My function"""
+        ...
+
+    schema = get_function_schema(func4)
+
+    assert schema["parameters"]["properties"]["animal"]["type"] == [
+        "string",
+        "integer",
+    ], "parameter animal should be a string"
+
+    assert schema["parameters"]["properties"]["animal"]["enum"] == [
+        "Cat",
+        "Dog",
+        1,
+    ], "parameter animal should have an enum"
+
+    assert (
+        schema["parameters"].get("required") is None
+        or "animal" not in schema["parameters"]["required"]
+    ), "animal should not be required"