Support PEP 604 Union Types (#347)

⚠️ Local Worker will not work with Slack toolkit until this is merged 

Support [PEP 604](https://peps.python.org/pep-0604/) Union Types for
tool input param types and return types.

https://github.com/ArcadeAI/arcade-ai/pull/345 introduced PEP 604 Union
Types into our toolkits, but the worker did not support this syntax, so
it was failing on start up if a tool used bar syntax in its return type
(Slack).

Additionally, optionals defined with `Union[T, None]` were already
supported, but didn't have any unit tests, so I added them.

arcade/arcade/core/catalog.py

@@ -48,6 +48,7 @@ from arcade.core.toolkit import Toolkit
 from arcade.core.utils import (
     does_function_return_value,
     first_or_none,
+    is_strict_optional,
     is_string_literal,
     is_union,
     snake_to_pascal_case,
@@ -481,10 +482,10 @@ def create_output_definition(func: Callable) -> ToolOutput:
         return_type = return_type.__origin__
 
     # Unwrap Optional types
-    is_optional = False
-    if get_origin(return_type) is Union and type(None) in get_args(return_type):
+    # Both Optional[T] and T | None are supported
+    is_optional = is_strict_optional(return_type)
+    if is_optional:
         return_type = next(arg for arg in get_args(return_type) if arg is not type(None))
-        is_optional = True
 
     wire_type_info = get_wire_type_info(return_type)
 
@@ -659,14 +660,9 @@ def extract_python_param_info(param: inspect.Parameter) -> ParamInfo:
 
     # Handle optional types
     # Both Optional[T] and T | None are supported
-    is_optional = False
-    if (
-        is_union(field_type)
-        and len(get_args(field_type)) == 2
-        and type(None) in get_args(field_type)
-    ):
+    is_optional = is_strict_optional(field_type)
+    if is_optional:
         field_type = next(arg for arg in get_args(field_type) if arg is not type(None))
-        is_optional = True
 
     # Union types are not currently supported
     # (other than optional, which is handled above)
@@ -703,10 +699,10 @@ def extract_pydantic_param_info(param: inspect.Parameter) -> ParamInfo:
     field_type = original_type
 
     # Unwrap Optional types
-    is_optional = False
-    if get_origin(field_type) is Union and type(None) in get_args(field_type):
+    # Both Optional[T] and T | None are supported
+    is_optional = is_strict_optional(field_type)
+    if is_optional:
         field_type = next(arg for arg in get_args(field_type) if arg is not type(None))
-        is_optional = True
 
     return ParamInfo(
         name=param.name,
@@ -732,7 +728,6 @@ def get_wire_type(
         float: "number",
         dict: "json",
     }
-
     outer_type_mapping: dict[type, WireType] = {
         list: "array",
         dict: "json",

arcade/arcade/core/utils.py

@@ -52,6 +52,14 @@ def is_union(_type: type) -> bool:
     return get_origin(_type) in {Union, UnionType}
 
 
+def is_strict_optional(_type: type) -> bool:
+    """
+    Returns True if the given type is a strict optional type, i.e. a union with exactly two types
+    where one type is None. This covers Optional[T], Union[T, None] and T | None
+    """
+    return is_union(_type) and len(get_args(_type)) == 2 and type(None) in get_args(_type)
+
+
 def does_function_return_value(func: Callable) -> bool:
     """
     Returns True if the given function returns a value, i.e. if it has a return statement with a value.

arcade/tests/core/test_catalog.py

@@ -94,14 +94,10 @@ def test_add_toolkit_type_error():
         mock_import.return_value = MagicMock()
         mock_getattr.return_value = InvalidTool()
 
-        # Assert that ToolDefinitionError is raised with the correct message
-        with pytest.raises(ToolDefinitionError) as exc_info:
+        # Assert that ToolDefinitionError is raised
+        with pytest.raises(ToolDefinitionError):
             catalog.add_toolkit(mock_toolkit)
 
-        assert "Type error encountered while adding tool invalid_tool from mock_module" in str(
-            exc_info.value
-        )
-
 
 def test_get_tool_by_name():
     catalog = ToolCatalog()

arcade/tests/core/utils/test_is_strict_optional.py

@@ -0,0 +1,26 @@
+from typing import Optional, Union
+
+import pytest
+
+from arcade.core.utils import is_strict_optional
+
+
+@pytest.mark.parametrize(
+    "type_input, expected",
+    [
+        (Union[int, None], True),
+        (Union[int, str, None], False),
+        (Union[int, str], False),
+        (Optional[int], True),
+        (int | None, True),
+        (None | int, True),
+        (int | str, False),
+        (int | str | None, False),
+        (int, False),
+        (str, False),
+        (list, False),
+        (dict, False),
+    ],
+)
+def test_is_optional_type(type_input, expected):
+    assert is_strict_optional(type_input) == expected

arcade/tests/tool/test_create_tool_definition.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Annotated, Literal, Optional
+from typing import Annotated, Literal, Optional, Union
 
 import pytest
 
@@ -172,12 +172,33 @@ def func_with_optional_param_with_default_value(
 
 
 @tool(desc="A function with an optional input parameter with bar syntax")
-def func_with_optional_param_with_bar_syntax(
+def func_with_optional_param_with_bar_syntax_1(
     param1: Annotated[str | None, "First param"] = None,
 ):
     pass
 
 
+@tool(desc="A function with an optional input parameter with bar syntax")
+def func_with_optional_param_with_bar_syntax_2(
+    param1: Annotated[None | str, "First param"] = None,
+):
+    pass
+
+
+@tool(desc="A function with an optional input parameter with union syntax")
+def func_with_optional_param_with_union_syntax_1(
+    param1: Annotated[Union[str, None], "First param"] = None,
+):
+    pass
+
+
+@tool(desc="A function with an optional input parameter with union syntax")
+def func_with_optional_param_with_union_syntax_2(
+    param1: Annotated[Union[None, str], "First param"] = None,
+):
+    pass
+
+
 @tool(desc="A function with multiple parameters, some with default values")
 def func_with_mixed_params(
     context: ToolContext,
@@ -240,6 +261,26 @@ def func_with_optional_return() -> Optional[str]:
     return "maybe output"
 
 
+@tool(desc="A function with an optional return type that uses bar syntax")
+def func_with_optional_return_with_bar_syntax_1() -> str | None:
+    return "maybe output"
+
+
+@tool(desc="A function with an optional return type that uses bar syntax")
+def func_with_optional_return_with_bar_syntax_2() -> None | str:
+    return "maybe output"
+
+
+@tool(desc="A function with an optional return type that uses union syntax")
+def func_with_optional_return_with_union_syntax_1() -> Union[str, None]:
+    return "maybe output"
+
+
+@tool(desc="A function with an optional return type that uses union syntax")
+def func_with_optional_return_with_union_syntax_2() -> Union[None, str]:
+    return "maybe output"
+
+
 @tool(desc="A function with a complex return type")
 def func_with_complex_return() -> dict[str, str]:
     return [{"key": "value"}]
@@ -536,7 +577,7 @@ def func_with_complex_return() -> dict[str, str]:
             id="func_with_optional_param_with_default_value",
         ),
         pytest.param(
-            func_with_optional_param_with_bar_syntax,
+            func_with_optional_param_with_bar_syntax_1,
             {
                 "input": ToolInput(
                     parameters=[
@@ -555,6 +596,57 @@ def func_with_complex_return() -> dict[str, str]:
             },
             id="func_with_optional_param_with_bar_syntax",
         ),
+        pytest.param(
+            func_with_optional_param_with_bar_syntax_2,
+            {
+                "input": ToolInput(
+                    parameters=[
+                        InputParameter(
+                            name="param1",
+                            description="First param",
+                            inferrable=True,
+                            required=False,
+                            value_schema=ValueSchema(val_type="string", enum=None),
+                        )
+                    ]
+                ),
+            },
+            id="func_with_optional_param_with_bar_syntax_2",
+        ),
+        pytest.param(
+            func_with_optional_param_with_union_syntax_1,
+            {
+                "input": ToolInput(
+                    parameters=[
+                        InputParameter(
+                            name="param1",
+                            description="First param",
+                            inferrable=True,
+                            required=False,
+                            value_schema=ValueSchema(val_type="string", enum=None),
+                        )
+                    ]
+                ),
+            },
+            id="func_with_optional_param_with_union_syntax_1",
+        ),
+        pytest.param(
+            func_with_optional_param_with_union_syntax_2,
+            {
+                "input": ToolInput(
+                    parameters=[
+                        InputParameter(
+                            name="param1",
+                            description="First param",
+                            inferrable=True,
+                            required=False,
+                            value_schema=ValueSchema(val_type="string", enum=None),
+                        )
+                    ]
+                ),
+            },
+            id="func_with_optional_param_with_union_syntax_2",
+        ),
         pytest.param(
             func_with_mixed_params,
             {
@@ -724,6 +816,54 @@ def func_with_complex_return() -> dict[str, str]:
             },
             id="func_with_optional_return",
         ),
+        pytest.param(
+            func_with_optional_return_with_bar_syntax_1,
+            {
+                "input": ToolInput(parameters=[]),
+                "output": ToolOutput(
+                    value_schema=ValueSchema(val_type="string", enum=None),
+                    available_modes=["value", "error", "null"],
+                    description="No description provided.",
+                ),
+            },
+            id="func_with_optional_return_with_bar_syntax_1",
+        ),
+        pytest.param(
+            func_with_optional_return_with_bar_syntax_2,
+            {
+                "input": ToolInput(parameters=[]),
+                "output": ToolOutput(
+                    value_schema=ValueSchema(val_type="string", enum=None),
+                    available_modes=["value", "error", "null"],
+                    description="No description provided.",
+                ),
+            },
+            id="func_with_optional_return_with_bar_syntax_2",
+        ),
+        pytest.param(
+            func_with_optional_return_with_union_syntax_1,
+            {
+                "input": ToolInput(parameters=[]),
+                "output": ToolOutput(
+                    value_schema=ValueSchema(val_type="string", enum=None),
+                    available_modes=["value", "error", "null"],
+                    description="No description provided.",
+                ),
+            },
+            id="func_with_optional_return_with_union_syntax_1",
+        ),
+        pytest.param(
+            func_with_optional_return_with_union_syntax_2,
+            {
+                "input": ToolInput(parameters=[]),
+                "output": ToolOutput(
+                    value_schema=ValueSchema(val_type="string", enum=None),
+                    available_modes=["value", "error", "null"],
+                    description="No description provided.",
+                ),
+            },
+            id="func_with_optional_return_with_union_syntax_2",
+        ),
         pytest.param(
             func_with_complex_return,
             {

arcade/tests/tool/test_create_tool_definition_errors.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Union
 
 import pytest
 
@@ -18,6 +18,16 @@ def func_with_missing_return_type():
     return "hello world"
 
 
+@tool(desc="A function with a union return type (illegal)")
+def func_with_union_return_type_1() -> Union[str, int]:
+    return "hello world"
+
+
+@tool(desc="A function with a union return type (illegal)")
+def func_with_union_return_type_2() -> str | int:
+    return "hello world"
+
+
 @tool(desc="A function with a parameter type (illegal)")
 def func_with_missing_param_type(param1):
     pass
@@ -34,7 +44,12 @@ def func_with_unsupported_param(param1: complex):
 
 
 @tool(desc="A function with a union parameter (illegal)")
-def func_with_union_param(param1: str | int):
+def func_with_union_param_1(param1: str | int):
+    pass
+
+
+@tool(desc="A function with a union parameter (illegal)")
+def func_with_union_param_2(param1: Union[str, int]):
     pass
 
 
@@ -95,9 +110,14 @@ def func_with_secret_requirement_invalid_type():
             id=func_with_unsupported_param.__name__,
         ),
         pytest.param(
-            func_with_union_param,
+            func_with_union_param_1,
             ToolDefinitionError,
-            id=func_with_union_param.__name__,
+            id=func_with_union_param_1.__name__,
+        ),
+        pytest.param(
+            func_with_union_param_2,
+            ToolDefinitionError,
+            id=func_with_union_param_2.__name__,
         ),
         pytest.param(
             func_with_multiple_context_params,
@@ -119,6 +139,16 @@ def func_with_secret_requirement_invalid_type():
             ToolDefinitionError,
             id=func_with_secret_requirement_invalid_type.__name__,
         ),
+        pytest.param(
+            func_with_union_return_type_1,
+            ToolDefinitionError,
+            id=func_with_union_return_type_1.__name__,
+        ),
+        pytest.param(
+            func_with_union_return_type_2,
+            ToolDefinitionError,
+            id=func_with_union_return_type_2.__name__,
+        ),
     ],
 )
 def test_missing_info_raises_error(func_under_test, exception_type):

arcade/tests/tool/test_create_tool_definition_pydantic.py

@@ -35,13 +35,27 @@ def func_takes_pydantic_field_with_description(
     return product_name
 
 
-@tool(desc="A function that accepts an Pydantic Field")
+@tool(desc="A function that accepts an optional Pydantic Field")
 def func_takes_pydantic_field_optional(
     product_name: Optional[str] = Field(None, description="The name of the product"),
 ) -> str:
     return product_name
 
 
+@tool(desc="A function that accepts an optional Pydantic Field with bar syntax")
+def func_takes_pydantic_field_optional_bar_syntax(
+    product_name: str | None = Field(None, description="The name of the product"),
+) -> str:
+    return product_name
+
+
+@tool(desc="A function that accepts an optional Pydantic Field with union syntax")
+def func_takes_pydantic_field_optional_union_syntax(
+    product_name: Union[str, None] = Field(None, description="The name of the product"),
+) -> str:
+    return product_name
+
+
 # Annotated[] takes precedence over Field() properties
 @tool(desc="A function that accepts an annotated Pydantic Field")
 def func_takes_pydantic_field_annotated_description(
@@ -174,6 +188,40 @@ def read_products(
             },
             id="func_takes_pydantic_field_optional",
         ),
+        pytest.param(
+            func_takes_pydantic_field_optional_bar_syntax,
+            {
+                "input": ToolInput(
+                    parameters=[
+                        InputParameter(
+                            name="product_name",
+                            description="The name of the product",
+                            required=False,
+                            inferrable=True,
+                            value_schema=ValueSchema(val_type="string", enum=None),
+                        )
+                    ]
+                )
+            },
+            id="func_takes_pydantic_field_optional_bar_syntax",
+        ),
+        pytest.param(
+            func_takes_pydantic_field_optional_union_syntax,
+            {
+                "input": ToolInput(
+                    parameters=[
+                        InputParameter(
+                            name="product_name",
+                            description="The name of the product",
+                            required=False,
+                            inferrable=True,
+                            value_schema=ValueSchema(val_type="string", enum=None),
+                        )
+                    ]
+                )
+            },
+            id="func_takes_pydantic_field_optional_union_syntax",
+        ),
         pytest.param(
             func_takes_pydantic_field_annotated_description,
             {

arcade/tests/tool/test_create_tool_definition_pydantic_errors.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Union
 
 import pytest
 from pydantic import Field
@@ -21,6 +21,20 @@ def field_with_literal_default_factory(
     pass
 
 
+@tool(desc="A function that accepts an optional Pydantic Field with non-strict optional syntax")
+def func_takes_pydantic_field_non_strict_optional_bar_syntax(
+    product_name: str | int | None = Field(None, description="The name of the product"),
+) -> str:
+    return product_name
+
+
+@tool(desc="A function that accepts an optional Pydantic Field with non-strict optional syntax")
+def func_takes_pydantic_field_non_strict_optional_union_syntax(
+    product_name: Union[str, int, None] = Field(None, description="The name of the product"),
+) -> str:
+    return product_name
+
+
 @pytest.mark.parametrize(
     "func_under_test, exception_type",
     [
@@ -29,6 +43,16 @@ def field_with_literal_default_factory(
             ToolDefinitionError,
             id=field_with_literal_default_factory.__name__,
         ),
+        pytest.param(
+            func_takes_pydantic_field_non_strict_optional_bar_syntax,
+            ToolDefinitionError,
+            id=func_takes_pydantic_field_non_strict_optional_bar_syntax.__name__,
+        ),
+        pytest.param(
+            func_takes_pydantic_field_non_strict_optional_union_syntax,
+            ToolDefinitionError,
+            id=func_takes_pydantic_field_non_strict_optional_union_syntax.__name__,
+        ),
     ],
 )
 def test_missing_info_raises_error(func_under_test, exception_type):