diff --git a/arraycontext/container/__init__.py b/arraycontext/container/__init__.py index 1046f990e2e48f09bff4e62130f0b2a067c417ae..ea208cb33fcb1c517baca9f6be0675ed614e522c 100644 --- a/arraycontext/container/__init__.py +++ b/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 diff --git a/arraycontext/container/traversal.py b/arraycontext/container/traversal.py index 76f7647179da16e41e321a1b1687864b15abbf01..ed6dc0c2482290a951258862aa392096a3040c16 100644 --- a/arraycontext/container/traversal.py +++ b/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)