Merge branch 'main' into cupyactx

Author: matthiasdiener

.github/workflows/ci.yml

@@ -52,12 +52,6 @@ jobs:
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
 
-                # NOTE: jax>=0.4.31 requires python 3.10 and uses pattern matching
-                # which conflicts with our mypy.python_version = '3.8' setting
-                CONDA_ENVIRONMENT=.test-conda-env-py3.yml
-                sed -i "s/jax/jax<0.4.31/" "$CONDA_ENVIRONMENT"
-                echo "- cupy" >> "$CONDA_ENVIRONMENT"
-
                 build_py_project_in_conda_env
                 python -m pip install mypy pytest
                 ./run-mypy.sh

.gitlab-ci.yml

@@ -107,9 +107,7 @@ Pylint:
 
 Mypy:
   script: |
-    # NOTE: jax>=0.4.31 requires python 3.10 and uses pattern matching
-    # which conflicts with our mypy.python_version = '3.8' setting
-    EXTRA_INSTALL="mypy pytest jax[cpu]<0.4.31"
+    EXTRA_INSTALL="mypy pytest"
 
     curl -L -O https://tiker.net/ci-support-v0
     . ./ci-support-v0

arraycontext/container/__init__.py

@@ -79,24 +79,14 @@
 THE SOFTWARE.
 """
 
+from collections.abc import Hashable, Sequence
 from functools import singledispatch
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Hashable,
-    Iterable,
-    Optional,
-    Protocol,
-    Sequence,
-    Tuple,
-    TypeVar,
-)
+from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, TypeVar
 
 # 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 TypeAlias
 
 from arraycontext.context import ArrayContext
 
@@ -154,8 +144,6 @@ class ArrayContainer(Protocol):
     # by dataclass_array_container, where it's used to recognize attributes
     # that are container-typed.
 
-    pass
-
 
 ArrayContainerT = TypeVar("ArrayContainerT", bound=ArrayContainer)
 
@@ -165,7 +153,7 @@ class NotAnArrayContainerError(TypeError):
 
 
 SerializationKey: TypeAlias = Hashable
-SerializedContainer: TypeAlias = Sequence[Tuple[SerializationKey, "ArrayOrContainer"]]
+SerializedContainer: TypeAlias = Sequence[tuple[SerializationKey, "ArrayOrContainer"]]
 
 
 @singledispatch
@@ -252,7 +240,7 @@ def is_array_container(ary: Any) -> bool:
 
 
 @singledispatch
-def get_container_context_opt(ary: ArrayContainer) -> Optional[ArrayContext]:
+def get_container_context_opt(ary: ArrayContainer) -> ArrayContext | None:
     """Retrieves the :class:`ArrayContext` from the container, if any.
 
     This function is not recursive, so it will only search at the root level
@@ -306,7 +294,7 @@ def _deserialize_ndarray_container(  # type: ignore[misc]
 # {{{ get_container_context_recursively
 
 def get_container_context_recursively_opt(
-        ary: ArrayContainer) -> Optional[ArrayContext]:
+        ary: ArrayContainer) -> ArrayContext | None:
     """Walks the :class:`ArrayContainer` hierarchy to find an
     :class:`ArrayContext` associated with it.
 
@@ -340,7 +328,7 @@ def get_container_context_recursively_opt(
         return actx
 
 
-def get_container_context_recursively(ary: ArrayContainer) -> Optional[ArrayContext]:
+def get_container_context_recursively(ary: ArrayContainer) -> ArrayContext | None:
     """Walks the :class:`ArrayContainer` hierarchy to find an
     :class:`ArrayContext` associated with it.
 
@@ -369,14 +357,16 @@ def get_container_context_recursively(ary: ArrayContainer) -> Optional[ArrayCont
 # FYI: This doesn't, and never should, make arraycontext directly depend on pymbolic.
 # (Though clearly there exists a dependency via loopy.)
 
-def _serialize_multivec_as_container(mv: MultiVector) -> Iterable[Tuple[Any, Any]]:
+def _serialize_multivec_as_container(mv: MultiVector) -> SerializedContainer:
     return list(mv.data.items())
 
 
-def _deserialize_multivec_as_container(template: MultiVector,
-        iterable: Iterable[Tuple[Any, Any]]) -> MultiVector:
+# FIXME: Ignored due to https://github.com/python/mypy/issues/13040
+def _deserialize_multivec_as_container(  # type: ignore[misc]
+        template: MultiVector,
+        serialized: SerializedContainer) -> MultiVector:
     from pymbolic.geometric_algebra import MultiVector
-    return MultiVector(dict(iterable), space=template.space)
+    return MultiVector(dict(serialized), space=template.space)
 
 
 def _get_container_context_opt_from_multivec(mv: MultiVector) -> None:

arraycontext/container/arithmetic.py

@@ -34,7 +34,8 @@
 """
 
 import enum
-from typing import Any, Callable, Optional, Tuple, TypeVar, Union
+from collections.abc import Callable
+from typing import Any, TypeVar
 from warnings import warn
 
 import numpy as np
@@ -90,7 +91,7 @@ class _OpClass(enum.Enum):
         ]
 
 
-def _format_unary_op_str(op_str: str, arg1: Union[Tuple[str, ...], str]) -> str:
+def _format_unary_op_str(op_str: str, arg1: tuple[str, ...] | str) -> str:
     if isinstance(arg1, tuple):
         arg1_entry, arg1_container = arg1
         return (f"{op_str.format(arg1_entry)} "
@@ -100,20 +101,14 @@ def _format_unary_op_str(op_str: str, arg1: Union[Tuple[str, ...], str]) -> str:
 
 
 def _format_binary_op_str(op_str: str,
-        arg1: Union[Tuple[str, str], str],
-        arg2: Union[Tuple[str, str], str]) -> str:
+        arg1: tuple[str, str] | str,
+        arg2: tuple[str, str] | str) -> str:
     if isinstance(arg1, tuple) and isinstance(arg2, tuple):
-        import sys
-        if sys.version_info >= (3, 10):
-            strict_arg = ", strict=__debug__"
-        else:
-            strict_arg = ""
-
         arg1_entry, arg1_container = arg1
         arg2_entry, arg2_container = arg2
         return (f"{op_str.format(arg1_entry, arg2_entry)} "
                 f"for {arg1_entry}, {arg2_entry} "
-                f"in zip({arg1_container}, {arg2_container}{strict_arg})")
+                f"in zip({arg1_container}, {arg2_container}, strict=__debug__)")
 
     elif isinstance(arg1, tuple):
         arg1_entry, arg1_container = arg1
@@ -160,23 +155,23 @@ class ComplainingNumpyNonObjectArray(metaclass=ComplainingNumpyNonObjectArrayMet
 
 def with_container_arithmetic(
             *,
-            number_bcasts_across: Optional[bool] = None,
-            bcasts_across_obj_array: Optional[bool] = None,
-            container_types_bcast_across: Optional[Tuple[type, ...]] = None,
+            number_bcasts_across: bool | None = None,
+            bcasts_across_obj_array: bool | None = None,
+            container_types_bcast_across: tuple[type, ...] | None = None,
             arithmetic: bool = True,
             matmul: bool = False,
             bitwise: bool = False,
             shift: bool = False,
-            _cls_has_array_context_attr: Optional[bool] = None,
-            eq_comparison: Optional[bool] = None,
-            rel_comparison: Optional[bool] = None,
+            _cls_has_array_context_attr: bool | None = None,
+            eq_comparison: bool | None = None,
+            rel_comparison: bool | None = None,
 
             # deprecated:
-            bcast_number: Optional[bool] = None,
-            bcast_obj_array: Optional[bool] = None,
+            bcast_number: bool | None = None,
+            bcast_obj_array: bool | None = None,
             bcast_numpy_array: bool = False,
-            _bcast_actx_array_type: Optional[bool] = None,
-            bcast_container_types: Optional[Tuple[type, ...]] = None,
+            _bcast_actx_array_type: bool | None = None,
+            bcast_container_types: tuple[type, ...] | None = None,
         ) -> Callable[[type], type]:
     """A class decorator that implements built-in operators for array containers
     by propagating the operations to the elements of the container.
@@ -482,7 +477,7 @@ def same_key(k1: T, k2: T) -> T:
             assert k1 == k2
             return k1
 
-        def tup_str(t: Tuple[str, ...]) -> str:
+        def tup_str(t: tuple[str, ...]) -> str:
             if not t:
                 return "()"
             else:
@@ -544,7 +539,8 @@ def {fname}(arg1):
                     _format_binary_op_str(op_str, expr_arg1, expr_arg2)
                     for (key_arg1, expr_arg1), (key_arg2, expr_arg2) in zip(
                         cls._serialize_init_arrays_code("arg1").items(),
-                        cls._serialize_init_arrays_code("arg2").items())
+                        cls._serialize_init_arrays_code("arg2").items(),
+                        strict=True)
                     })
             bcast_init_args_arg1_is_outer = cls._deserialize_init_arrays_code("arg1", {
                     key_arg1: _format_binary_op_str(op_str, expr_arg1, "arg2")

arraycontext/container/dataclass.py

@@ -31,7 +31,7 @@
 """
 
 from dataclasses import Field, fields, is_dataclass
-from typing import Tuple, Union, get_args, get_origin
+from typing import Union, get_args, get_origin
 
 from arraycontext.container import is_array_container_type
 
@@ -83,14 +83,15 @@ def is_array_field(f: Field) -> bool:
                         f"Field '{f.name}' union contains non-array container "
                         "arguments. All arguments must be array containers.")
 
+        if isinstance(f.type, str):
+            raise TypeError(
+                f"String annotation on field '{f.name}' not supported. "
+                "(this may be due to 'from __future__ import annotations')")
+
         if __debug__:
             if not f.init:
                 raise ValueError(
-                        f"'init=False' field not allowed: '{f.name}'")
-
-            if isinstance(f.type, str):
-                raise TypeError(
-                        f"string annotation on field '{f.name}' not supported")
+                        f"Field with 'init=False' not allowed: '{f.name}'")
 
             # NOTE:
             # * `_BaseGenericAlias` catches `List`, `Tuple`, etc.
@@ -99,15 +100,15 @@ def is_array_field(f: Field) -> bool:
                 _BaseGenericAlias,
                 _SpecialForm,
             )
-            if isinstance(f.type, (_BaseGenericAlias, _SpecialForm)):
+            if isinstance(f.type, _BaseGenericAlias | _SpecialForm):
                 # NOTE: anything except a Union is not allowed
                 raise TypeError(
-                        f"typing annotation not supported on field '{f.name}': "
+                        f"Typing annotation not supported on field '{f.name}': "
                         f"'{f.type!r}'")
 
             if not isinstance(f.type, type):
                 raise TypeError(
-                        f"field '{f.name}' not an instance of 'type': "
+                        f"Field '{f.name}' not an instance of 'type': "
                         f"'{f.type!r}'")
 
         return is_array_type(f.type)
@@ -124,8 +125,8 @@ def is_array_field(f: Field) -> bool:
 
 def inject_dataclass_serialization(
         cls: type,
-        array_fields: Tuple[Field, ...],
-        non_array_fields: Tuple[Field, ...]) -> type:
+        array_fields: tuple[Field, ...],
+        non_array_fields: tuple[Field, ...]) -> type:
     """Implements :func:`~arraycontext.serialize_container` and
     :func:`~arraycontext.deserialize_container` for the given dataclass *cls*.