Add all_type_leaves_satisfy_predicate, refactor type tests to use it

Author: inducer

arraycontext/container/__init__.py

@@ -135,7 +135,6 @@
 from typing import (
     TYPE_CHECKING,
     TypeAlias,
-    get_origin,
 )
 
 # For use in singledispatch type annotations, because sphinx can't figure out
@@ -147,14 +146,15 @@
 from pytools.obj_array import ObjectArray, ObjectArrayND as ObjectArrayND
 
 from arraycontext.typing import (
-    ArithArrayContainer,
+    ArithArrayContainer as ArithArrayContainer,
     ArrayContainer,
     ArrayContainerT,
     ArrayOrArithContainer,
     ArrayOrArithContainerOrScalar as ArrayOrArithContainerOrScalar,
     ArrayOrContainerOrScalar,
     _UserDefinedArithArrayContainer,
     _UserDefinedArrayContainer,
+    all_type_leaves_satisfy_predicate,
 )
 
 
@@ -233,23 +233,15 @@ def is_array_container_type(cls: type | GenericAlias | UnionType) -> bool:
         function will say that :class:`numpy.ndarray` is an array container
         type, only object arrays *actually are* array containers.
     """
-    if cls is ArrayContainer or cls is ArithArrayContainer:
-        return True
+    def pred(tp: type) -> bool:
+        return (
+                tp is ObjectArray
+                or tp is _UserDefinedArrayContainer
+                or tp is _UserDefinedArithArrayContainer
+                or (serialize_container.dispatch(tp)
+                    is not serialize_container.__wrapped__))  # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
 
-    origin = get_origin(cls)
-    if origin is not None:
-        cls = origin  # pyright: ignore[reportAny]
-
-    assert isinstance(cls, type), (
-        f"must pass a {type!r}, not a '{cls!r}'")
-
-    return (
-            cls is ObjectArray
-            or cls is ArrayContainer  # pyright: ignore[reportUnnecessaryComparison]
-            or cls is _UserDefinedArrayContainer
-            or cls is _UserDefinedArithArrayContainer
-            or (serialize_container.dispatch(cls)
-                is not serialize_container.__wrapped__))  # type:ignore[attr-defined]
+    return all_type_leaves_satisfy_predicate(pred, cls)
 
 
 def is_array_container(ary: object) -> TypeIs[ArrayContainer]:
@@ -265,7 +257,7 @@ def is_array_container(ary: object) -> TypeIs[ArrayContainer]:
             "cheaper option, see is_array_container_type.",
             DeprecationWarning, stacklevel=2)
     return (serialize_container.dispatch(ary.__class__)
-            is not serialize_container.__wrapped__       # type:ignore[attr-defined]
+            is not serialize_container.__wrapped__       # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
             # numpy values with scalar elements aren't array containers
             and not (isinstance(ary, np.ndarray)
                      and ary.dtype.kind != "O")

arraycontext/container/dataclass.py

@@ -37,30 +37,19 @@
 THE SOFTWARE.
 """
 
-# The import of 'Union' is type-ignored below because we're specifically importing
-# Union to pick apart type annotations.
-
 from dataclasses import fields, is_dataclass
 from typing import (
     TYPE_CHECKING,
     NamedTuple,
     TypeVar,
-    Union,  # pyright: ignore[reportDeprecated]
-    cast,
-    get_args,
-    get_origin,
 )
 from warnings import warn
 
 import numpy as np
 
-from pytools.obj_array import ObjectArray
-
 from arraycontext.container import is_array_container_type
 from arraycontext.typing import (
-    ArrayContainer,
-    ArrayOrContainer,
-    ArrayOrContainerOrScalar,
+    all_type_leaves_satisfy_predicate,
     is_scalar_type,
 )
 
@@ -83,17 +72,30 @@ class _Field(NamedTuple):
     type: type
 
 
-def _is_array_or_container_type(tp: type | GenericAlias | UnionType, /) -> bool:
-    if tp is np.ndarray:
-        warn("Encountered 'numpy.ndarray' in a dataclass_array_container. "
-             "This is deprecated and will stop working in 2026. "
-             "If you meant an object array, use pytools.obj_array.ObjectArray. "
-             "For other uses, file an issue to discuss.",
-             DeprecationWarning, stacklevel=3)
-        return True
+def _is_array_or_container_type(
+            tp: type | GenericAlias | UnionType | TypeVar, /, *,
+            allow_scalar: bool = True,
+            require_homogeneity: bool = True,
+        ) -> bool:
+    def _is_array_or_container_or_scalar(tp: type) -> bool:
+        if tp is np.ndarray:
+            warn("Encountered 'numpy.ndarray' in a dataclass_array_container. "
+                 "This is deprecated and will stop working in 2026. "
+                 "If you meant an object array, use pytools.obj_array.ObjectArray. "
+                 "For other uses, file an issue to discuss.",
+                 DeprecationWarning, stacklevel=1)
+            return True
+
+        from arraycontext import Array
+
+        return (
+                is_array_container_type(tp)
+                or tp is Array
+                or (allow_scalar and is_scalar_type(tp)))
 
-    from arraycontext import Array
-    return tp is Array or is_array_container_type(tp)
+    return all_type_leaves_satisfy_predicate(
+                    _is_array_or_container_or_scalar, tp,
+                    require_homogeneity=require_homogeneity)
 
 
 def dataclass_array_container(cls: type[T]) -> type[T]:
@@ -120,8 +122,6 @@ def dataclass_array_container(cls: type[T]) -> type[T]:
         means that *cls* must live in a module that is importable.
     """
 
-    from types import GenericAlias, UnionType
-
     assert is_dataclass(cls)
 
     def is_array_field(f: _Field) -> bool:
@@ -139,61 +139,17 @@ def is_array_field(f: _Field) -> bool:
         #
         # This is not set in stone, but mostly driven by current usage!
 
-        # pyright has no idea what we're up to. :)
-        if field_type is ArrayContainer:  # pyright: ignore[reportUnnecessaryComparison]
-            return True
-        if field_type is ArrayOrContainer:  # pyright: ignore[reportUnnecessaryComparison]
-            return True
-        if field_type is ArrayOrContainerOrScalar:  # pyright: ignore[reportUnnecessaryComparison]
-            return True
-
-        origin = get_origin(field_type)
-
-        if origin is ObjectArray:
-            return True
-
-        # NOTE: `UnionType` is returned when using `Type1 | Type2`
-        if origin in (Union, UnionType):  # pyright: ignore[reportDeprecated]
-            for arg in get_args(field_type):  # pyright: ignore[reportAny]
-                if not (
-                        _is_array_or_container_type(cast("type", arg))
-                        or is_scalar_type(cast("type", arg))):
-                    raise TypeError(
-                            f"Field '{f.name}' union contains non-array container "
-                            f"type '{arg}'. All types must be array containers "
-                            "or arrays or scalars."
-                        )
-
-            return True
-
         # NOTE: this should never happen due to using `inspect.get_annotations`
         assert not isinstance(field_type, str)
 
-        if __debug__:
-            if not f.init:
-                raise ValueError(
-                        f"Field with 'init=False' not allowed: '{f.name}'")
-
-            # NOTE:
-            # * `GenericAlias` catches typed `list`, `tuple`, etc.
-            # * `_BaseGenericAlias` catches `List`, `Tuple`, etc.
-            # * `_SpecialForm` catches `Any`, `Literal`, etc.
-            from typing import (  # type: ignore[attr-defined]
-                _BaseGenericAlias,
-                _SpecialForm,
-            )
-            if isinstance(field_type, GenericAlias | _BaseGenericAlias | _SpecialForm):
-                # NOTE: anything except a Union is not allowed
-                raise TypeError(
-                        f"Type annotation not supported on field '{f.name}': "
-                        f"'{field_type!r}'")
-
-            if not isinstance(field_type, type):
-                raise TypeError(
-                        f"Field '{f.name}' not an instance of 'type': "
-                        f"'{field_type!r}'")
-
-        return _is_array_or_container_type(field_type)
+        if not f.init:
+            raise ValueError(
+                    f"Field with 'init=False' not allowed: '{f.name}'")
+
+        try:
+            return _is_array_or_container_type(field_type)
+        except TypeError as e:
+            raise TypeError(f"Field '{f.name}': {e}") from None
 
     from pytools import partition
 

arraycontext/typing.py

@@ -71,6 +71,11 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
+# The import of 'Union' is type-ignored below because we're specifically importing
+# Union to pick apart old/deprecated type annotations.
+
+from functools import partial
+from types import GenericAlias, UnionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -80,7 +85,9 @@
     SupportsInt,
     TypeAlias,
     TypeVar,
+    Union,  # pyright: ignore[reportDeprecated]
     cast,
+    get_args,
     get_origin,
     overload,
 )
@@ -89,10 +96,13 @@
 from typing_extensions import Self, TypeIs
 
 from pymbolic.typing import Integer, Scalar as _Scalar
+from pytools import partition2
 from pytools.obj_array import ObjectArrayND
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from numpy.typing import DTypeLike
 
     from pymbolic.typing import Integer
@@ -296,3 +306,56 @@ def shape_is_int_only(shape: tuple[Array | Integer, ...], /) -> tuple[int, ...]:
                 ) from None
 
     return tuple(res)
+
+
+def all_type_leaves_satisfy_predicate(
+            predicate: Callable[[type], bool],
+            tp: type | GenericAlias | UnionType | TypeVar,
+            /, *,
+            require_homogeneity: bool = False,
+            allow_containers_with_satisfying_types: bool = False,
+        ) -> bool:
+    # This is horrible and brittle. I'm sorry.
+
+    rec = partial(
+                  all_type_leaves_satisfy_predicate,
+                  predicate,
+                  require_homogeneity=require_homogeneity,
+                  allow_containers_with_satisfying_types=allow_containers_with_satisfying_types
+              )
+    origin = get_origin(tp)
+    args = get_args(tp)
+    tp_or_origin = tp if origin is None else origin
+
+    if isinstance(tp_or_origin, TypeVar):
+        bound = cast("type | None", tp_or_origin.__bound__)
+        if bound is None:
+            return False
+        else:
+            return rec(bound)
+
+    # NOTE: `UnionType` is returned when using `Type1 | Type2`
+    if origin in (Union, UnionType):  # pyright: ignore[reportDeprecated]
+        yes_types, no_types = partition2(
+            (rec(arg), arg) for arg in args)  # pyright: ignore[reportAny]
+        if require_homogeneity and yes_types and no_types:
+            raise TypeError(f"union '{tp}' is non-homogeneous "
+                            f"in whether it satisfies '{predicate}'")
+
+        return not no_types
+
+    if not isinstance(tp_or_origin, type):
+        raise TypeError(f"encountered non-type '{type(tp_or_origin)!r}'")
+
+    if predicate(tp_or_origin):
+        return True
+
+    if args and not allow_containers_with_satisfying_types:
+        # assume these are containers
+        has_sat_types = any(rec(arg) for arg in args)  # pyright: ignore[reportAny]
+
+        if has_sat_types:
+            raise TypeError(f"container '{tp}' has an element type "
+                            f"satisfying '{predicate}'")
+
+    return False

test/test_utils.py

@@ -72,9 +72,9 @@ def test_dataclass_array_container() -> None:
     class ArrayContainerWithOptional:
         x: np.ndarray
         # Deliberately left as Optional to test compatibility.
-        y: Optional[np.ndarray]  # noqa: UP045
+        y: Optional[np.ndarray]  # noqa: UP045  # pyright: ignore[reportDeprecated]
 
-    with pytest.raises(TypeError, match="Field 'y' union contains non-array"):
+    with pytest.raises(TypeError, match=r"Field 'y':.*non-homogeneous.*"):
         # NOTE: cannot have wrapped annotations (here by `Optional`)
         dataclass_array_container(ArrayContainerWithOptional)
 
@@ -88,15 +88,15 @@ class ArrayContainerWithTuple:
         # Deliberately left as Tuple to test compatibility.
         y: Tuple[Array, Array]  # noqa: UP006
 
-    with pytest.raises(TypeError, match="Type annotation not supported on field 'y'"):
+    with pytest.raises(TypeError, match=r"Field 'y':.*has an element type.*"):
         dataclass_array_container(ArrayContainerWithTuple)
 
     @dataclass
     class ArrayContainerWithTupleAlt:
         x: Array
         y: tuple[Array, Array]
 
-    with pytest.raises(TypeError, match="Type annotation not supported on field 'y'"):
+    with pytest.raises(TypeError, match=r"Field 'y':.*has an element type.*"):
         dataclass_array_container(ArrayContainerWithTupleAlt)
 
     # }}}
@@ -159,12 +159,21 @@ class ArrayContainerWithUnionAlt:
     @dataclass
     class ArrayContainerWithWrongUnion:
         x: np.ndarray
-        y: np.ndarray | list[bool]
+        y: np.ndarray | list[str]
 
-    with pytest.raises(TypeError, match="Field 'y' union contains non-array container"):
-        # NOTE: bool is not an ArrayContainer, so y should fail
+    with pytest.raises(TypeError, match=r"Field 'y':.*non-homogeneous.*"):
+        # NOTE: str is not an ArrayContainer, so y should fail
         dataclass_array_container(ArrayContainerWithWrongUnion)
 
+    @dataclass
+    class ArrayContainerWithWrongUnion2:
+        x: np.ndarray
+        y: np.ndarray | str
+
+    with pytest.raises(TypeError, match=r"Field 'y':.*non-homogeneous.*"):
+        # NOTE: str is not an ArrayContainer, so y should fail
+        dataclass_array_container(ArrayContainerWithWrongUnion2)
+
     # }}}
 
     # {{{ optional union
@@ -174,7 +183,7 @@ class ArrayContainerWithOptionalUnion:
         x: np.ndarray
         y: np.ndarray | None
 
-    with pytest.raises(TypeError, match="Field 'y' union contains non-array container"):
+    with pytest.raises(TypeError, match=r"Field 'y':.*non-homogeneous.*"):
         # NOTE: None is not an ArrayContainer, so y should fail
         dataclass_array_container(ArrayContainerWithWrongUnion)