add ArrayT to fix mypy without limiting ContainerT usefulness

arraycontext/container/__init__.py

@@ -3,6 +3,11 @@
 """
 .. currentmodule:: arraycontext
 
+.. class:: ArrayT
+    :canonical: arraycontext.container.ArrayT
+
+    :class:`~typing.TypeVar` for arrays.
+
 .. class:: ContainerT
     :canonical: arraycontext.container.ContainerT
 
@@ -54,11 +59,12 @@ THE SOFTWARE.
 
 from functools import singledispatch
 from arraycontext.context import ArrayContext
-from typing import Any, Iterable, Tuple, TypeVar, Optional
+from typing import Any, Iterable, Tuple, TypeVar, Optional, Union
 import numpy as np
 
+ArrayT = TypeVar("ArrayT")
 ContainerT = TypeVar("ContainerT")
-ArrayOrContainerT = TypeVar("ArrayOrContainerT")
+ArrayOrContainerT = Union[ArrayT, ContainerT]
 
 
 # {{{ ArrayContainer

arraycontext/container/traversal.py

@@ -142,13 +142,13 @@ def _multimap_array_container_impl(
     if len(container_indices) == 1 and reduce_func is None:
         # NOTE: if we just have one ArrayContainer in args, passing it through
         # _map_array_container_impl should be faster
-        def wrapper(ary: ArrayOrContainerT) -> ArrayOrContainerT:
+        def wrapper(ary: ContainerT) -> ContainerT:
             new_args = list(args)
             new_args[container_indices[0]] = ary
             return f(*new_args)
 
         update_wrapper(wrapper, f)
-        template_ary: ArrayOrContainerT = args[container_indices[0]]
+        template_ary: ContainerT = args[container_indices[0]]
         return _map_array_container_impl(
                 wrapper, template_ary,
                 leaf_cls=leaf_cls, recursive=recursive)