Centralize typing stuff in arraycontext.typing

Author: inducer

arraycontext/__init__.py

@@ -30,9 +30,6 @@
 """
 
 from .container import (
-    ArithArrayContainer,
-    ArrayContainer,
-    ArrayContainerT,
     NotAnArrayContainerError,
     SerializationKey,
     SerializedContainer,
@@ -72,9 +69,26 @@
     with_array_context,
 )
 from .context import (
-    Array,
     ArrayContext,
     ArrayContextFactory,
+    tag_axes,
+)
+from .impl.jax import EagerJAXArrayContext
+from .impl.numpy import NumpyArrayContext
+from .impl.pyopencl import PyOpenCLArrayContext
+from .impl.pytato import PytatoJAXArrayContext, PytatoPyOpenCLArrayContext
+from .loopy import make_loopy_program
+from .pytest import (
+    PytestArrayContextFactory,
+    PytestPyOpenCLArrayContextFactory,
+    pytest_generate_tests_for_array_contexts,
+)
+from .transform_metadata import CommonSubexpressionTag, ElementwiseMapKernelTag
+from .typing import (
+    ArithArrayContainer,
+    Array,
+    ArrayContainer,
+    ArrayContainerT,
     ArrayOrArithContainer,
     ArrayOrArithContainerOrScalar,
     ArrayOrArithContainerOrScalarT,
@@ -86,21 +100,10 @@
     ArrayOrScalar,
     ArrayOrScalarT,
     ArrayT,
+    ContainerOrScalarT,
     Scalar,
     ScalarLike,
-    tag_axes,
-)
-from .impl.jax import EagerJAXArrayContext
-from .impl.numpy import NumpyArrayContext
-from .impl.pyopencl import PyOpenCLArrayContext
-from .impl.pytato import PytatoJAXArrayContext, PytatoPyOpenCLArrayContext
-from .loopy import make_loopy_program
-from .pytest import (
-    PytestArrayContextFactory,
-    PytestPyOpenCLArrayContextFactory,
-    pytest_generate_tests_for_array_contexts,
 )
-from .transform_metadata import CommonSubexpressionTag, ElementwiseMapKernelTag
 
 
 __all__ = (
@@ -123,6 +126,7 @@
     "ArrayT",
     "BcastUntilActxArray",
     "CommonSubexpressionTag",
+    "ContainerOrScalarT",
     "EagerJAXArrayContext",
     "ElementwiseMapKernelTag",
     "NotAnArrayContainerError",

arraycontext/container/__init__.py

@@ -49,9 +49,7 @@
         This should be considered experimental for now, and it may well change.
 
 .. autoclass:: ArithArrayContainer
-.. class:: ArrayContainerT
-
-    A type variable with a lower bound of :class:`ArrayContainer`.
+.. autoclass:: ArrayContainerT
 
 .. autoexception:: NotAnArrayContainerError
 
@@ -125,80 +123,35 @@
 from types import GenericAlias, UnionType
 from typing import (
     TYPE_CHECKING,
-    ClassVar,
-    Protocol,
     TypeAlias,
-    TypeVar,
     get_origin,
 )
 
 # For use in singledispatch type annotations, because sphinx can't figure out
 # what 'np' is.
 import numpy
 import numpy as np
-from typing_extensions import Self, TypeIs
+from typing_extensions import TypeIs
 
-from pytools.obj_array import ObjectArrayND
+from pytools.obj_array import ObjectArrayND as ObjectArrayND
 
-from arraycontext.context import (
+from arraycontext.typing import (
+    ArrayContainer,
+    ArrayContainerT,
     ArrayOrArithContainer,
-    ArrayOrArithContainerOrScalar,
+    ArrayOrArithContainerOrScalar as ArrayOrArithContainerOrScalar,
     ArrayOrContainerOrScalar,
 )
 
 
 if TYPE_CHECKING:
     from pymbolic.geometric_algebra import CoeffT, MultiVector
 
-    from arraycontext.context import ArrayContext, ArrayOrScalar
-
-
-# {{{ ArrayContainer
-
-class _UserDefinedArrayContainer(Protocol):
-    # This is used as a type annotation in dataclasses that are processed
-    # by dataclass_array_container, where it's used to recognize attributes
-    # that are container-typed.
-
-    # This method prevents ArrayContainer from matching any object, while
-    # matching numpy object arrays and many array containers.
-    __array_ufunc__: ClassVar[None]
-
-
-ArrayContainer: TypeAlias = (
-    ObjectArrayND[ArrayOrContainerOrScalar]
-    | _UserDefinedArrayContainer
-    )
-
-
-class _UserDefinedArithArrayContainer(_UserDefinedArrayContainer, Protocol):
-    # This is loose and permissive, assuming that any array can be added
-    # to any container. The alternative would be to plaster type-ignores
-    # on all those uses. Achieving typing precision on what broadcasting is
-    # allowable seems like a huge endeavor and is likely not feasible without
-    # a mypy plugin. Maybe some day? -AK, November 2024
-
-    def __neg__(self) -> Self: ...
-    def __abs__(self) -> Self: ...
-    def __add__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __radd__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __sub__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __rsub__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __mul__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __rmul__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __truediv__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __rtruediv__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __pow__(self, other: ArrayOrScalar | Self) -> Self: ...
-    def __rpow__(self, other: ArrayOrScalar | Self) -> Self: ...
-
-
-ArithArrayContainer: TypeAlias = (
-    ObjectArrayND[ArrayOrArithContainerOrScalar]
-    | _UserDefinedArithArrayContainer)
-
+    from arraycontext.context import ArrayContext
+    from arraycontext.typing import ArrayOrScalar as ArrayOrScalar
 
-ArrayContainerT = TypeVar("ArrayContainerT", bound=ArrayContainer)
 
+# {{{ ArrayContainer traversals
 
 class NotAnArrayContainerError(TypeError):
     """:class:`TypeError` subclass raised when an array container is expected."""

arraycontext/container/arithmetic.py

@@ -68,8 +68,8 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from arraycontext.context import (
-        ArrayContext,
+    from arraycontext.context import ArrayContext
+    from arraycontext.typing import (
         ArrayOrContainer,
         ArrayOrContainerOrScalar,
     )

arraycontext/container/dataclass.py

@@ -56,8 +56,12 @@
 
 from pytools.obj_array import ObjectArray
 
-from arraycontext.container import ArrayContainer, is_array_container_type
-from arraycontext.context import ArrayOrContainer, ArrayOrContainerOrScalar
+from arraycontext.container import is_array_container_type
+from arraycontext.typing import (
+    ArrayContainer,
+    ArrayOrContainer,
+    ArrayOrContainerOrScalar,
+)
 
 
 if TYPE_CHECKING:

arraycontext/container/traversal.py

@@ -33,14 +33,26 @@
 .. autofunction:: unflatten
 .. autofunction:: flat_size_and_dtype
 
-Numpy conversion
-~~~~~~~~~~~~~~~~
-.. autofunction:: from_numpy
-.. autofunction:: to_numpy
-
 Algebraic operations
 ~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: outer
+
+.. currentmodule:: arraycontext.traversal
+
+References
+----------
+
+.. class:: ArrayOrScalar
+
+    See :class:`arraycontext.ArrayOrScalar`.
+
+.. class:: ArrayOrContainer
+
+    See :class:`arraycontext.ArrayOrContainer`.
+
+.. class:: ArrayContainerT
+
+    See :class:`arraycontext.ArrayContainerT`.
 """
 
 from __future__ import annotations
@@ -84,24 +96,27 @@
 )
 
 from arraycontext.container import (
-    ArrayContainer,
-    ArrayContainerT,
     NotAnArrayContainerError,
     SerializationKey,
     deserialize_container,
     get_container_context_recursively_opt,
     is_array_container,
     serialize_container,
 )
-from arraycontext.context import is_scalar_like, shape_is_int_only
+from arraycontext.typing import (
+    ArrayContainer,
+    ArrayContainerT,
+    is_scalar_like,
+    shape_is_int_only,
+)
 
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Collection, Iterable
 
-    from arraycontext.context import (
+    from arraycontext.context import ArrayContext
+    from arraycontext.typing import (
         Array,
-        ArrayContext,
         ArrayOrContainer,
         ArrayOrContainerOrScalar,
         ArrayOrContainerOrScalarT,