feat: add tests for union/optional types

Author: msto

dataclass_io/_lib/typing_extensions.py

@@ -0,0 +1,38 @@
+from types import UnionType
+from typing import Union
+from typing import get_args
+from typing import get_origin
+
+NoneType = type(None)
+"""Helpful alias for `type(None)`."""
+
+
+def is_union(type_: type) -> bool:
+    """
+    True if `type_` is a union type.
+
+    When declared with `Union[T, ...]` or `Optional[T]`, `get_origin()` returns `typing.Union`.
+    When declared with PEP604 syntax `T | ...`, `get_origin()` returns `types.UnionType`.
+
+    Args:
+        type_: The type to check.
+
+    Returns:
+        True if `type_` is a union type.
+    """
+    return get_origin(type_) is Union or get_origin(type_) is UnionType
+
+
+def is_optional(type_: type) -> bool:
+    """
+    True if `_type` is `Optional`, or the union of a single type and `None`.
+
+    Args:
+        type_: The type to check.
+
+    Returns:
+        True if `_type` is `Optional[T]`, `Union[T, None]` or `T | None`.
+    """
+    type_args: tuple[type] = get_args(type_)
+
+    return is_union(type_) and (NoneType in type_args) and (len(type_args) == 2)

tests/_lib/test_typing_extensions

@@ -0,0 +1,32 @@
+from typing import Union, Optional, Any
+from dataclass_io._lib.typing_extensions import is_optional
+from dataclass_io._lib.typing_extensions import is_union
+
+
+def test_is_union() -> None:
+    """Test that we can identify a union type."""
+
+    assert is_union(Union[int, None])
+    assert is_union(Union[int, str])
+    assert is_union(Union[int, str, float])
+    assert is_union(Optional[int])
+
+    assert not is_union(int)
+    assert not is_union(str)
+    assert not is_union(dict[int, str])
+
+
+def test_is_optional() -> None:
+    """Test that we can identify an Optional type."""
+
+    assert is_optional(Union[int, None])
+    assert is_optional(int | None)
+    assert is_optional(Optional[int])
+
+    assert is_optional(Union[dict[str, Any], None])
+    assert is_optional(dict[str, Any] | None)
+    assert is_optional(Optional[dict[str, Any]])
+
+    assert not is_optional(int)
+    assert not is_optional(Union[int, str])
+    assert not is_optional(dict[str, Any])