bpr things

Author: matthiasdiener

arraycontext/container/__init__.py

@@ -94,6 +94,8 @@
 
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from pymbolic.geometric_algebra import MultiVector
 
     from arraycontext import ArrayOrContainer
@@ -283,7 +285,7 @@ def get_container_context_opt(ary: ArrayContainer) -> ArrayContext | None:
 
 @serialize_container.register(np.ndarray)
 def _serialize_ndarray_container(
-        ary: numpy.ndarray) -> SerializedContainer:
+        ary: numpy.ndarray[Any, Any]) -> SerializedContainer:
     if ary.dtype.char != "O":
         raise NotAnArrayContainerError(
                 f"cannot serialize '{type(ary).__name__}' with dtype '{ary.dtype}'")
@@ -303,8 +305,8 @@ def _serialize_ndarray_container(
 @deserialize_container.register(np.ndarray)
 # https://github.com/python/mypy/issues/13040
 def _deserialize_ndarray_container(  # type: ignore[misc]
-        template: numpy.ndarray,
-        serialized: SerializedContainer) -> numpy.ndarray:
+        template: numpy.ndarray[Any, Any],
+        serialized: SerializedContainer) -> numpy.ndarray[Any, Any]:
     # disallow subclasses
     assert type(template) is np.ndarray
     assert template.dtype.char == "O"

arraycontext/container/traversal.py

@@ -917,7 +917,7 @@ def _flat_size(subary: ArrayOrContainer) -> Array | Integer:
 # {{{ numpy conversion
 
 def from_numpy(
-        ary: np.ndarray | ScalarLike,
+        ary: np.ndarray[Any, Any] | ScalarLike,
         actx: ArrayContext) -> ArrayOrContainerOrScalar:
     """Convert all :mod:`numpy` arrays in the :class:`~arraycontext.ArrayContainer`
     to the base array type of :class:`~arraycontext.ArrayContext`.

arraycontext/context.py

@@ -282,7 +282,7 @@ def __rtruediv__(self, other: Self | ScalarLike) -> Array: ...
 ContainerOrScalarT = TypeVar("ContainerOrScalarT", bound="ArrayContainer | ScalarLike")
 
 
-NumpyOrContainerOrScalar = Union[np.ndarray, "ArrayContainer", ScalarLike]
+NumpyOrContainerOrScalar = Union[np.ndarray[Any, Any], "ArrayContainer", ScalarLike]
 
 # }}}
 
@@ -358,7 +358,7 @@ def zeros(self,
         return self.np.zeros(shape, dtype)
 
     @overload
-    def from_numpy(self, array: np.ndarray) -> Array:
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
         ...
 
     @overload
@@ -379,7 +379,7 @@ def from_numpy(self,
         """
 
     @overload
-    def to_numpy(self, array: Array) -> np.ndarray:
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
         ...
 
     @overload

arraycontext/impl/cupy/__init__.py

@@ -76,7 +76,7 @@ def clone(self):
         return type(self)()
 
     @overload
-    def from_numpy(self, array: np.ndarray) -> Array:
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
         ...
 
     @overload
@@ -91,7 +91,7 @@ def from_numpy(self,
                                   actx=self)
 
     @overload
-    def to_numpy(self, array: Array) -> np.ndarray:
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
         ...
 
     @overload

arraycontext/impl/cupy/fake_numpy.py

@@ -178,7 +178,7 @@ def linspace(self, *args, **kwargs):
         import cupy as cp
         return cp.linspace(*args, **kwargs)
 
-    def zeros_like(self, ary):
+    def zeros_like(self, ary):  # pyright: ignore[reportIncompatibleMethodOverride]
         import cupy as cp
         if isinstance(ary, int | float | complex):
             # Cupy does not support zeros_like with scalar arguments

arraycontext/impl/numpy/__init__.py

@@ -86,7 +86,7 @@ def clone(self):
         return type(self)()
 
     @overload
-    def from_numpy(self, array: np.ndarray) -> Array:
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
         ...
 
     @overload
@@ -99,7 +99,7 @@ def from_numpy(self,
         return array
 
     @overload
-    def to_numpy(self, array: Array) -> np.ndarray:
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
         ...
 
     @overload

arraycontext/impl/numpy/fake_numpy.py

@@ -150,7 +150,7 @@ def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
                 return false_ary
             return np.logical_and.reduce(
                     [(true_ary if kx_i == ky_i else false_ary)
-                        and cast(np.ndarray, self.array_equal(x_i, y_i))
+                        and cast(np.ndarray, self.array_equal(x_i, y_i))  # pyright: ignore[reportMissingTypeArgument]
                         for (kx_i, x_i), (ky_i, y_i)
                         in zip(serialized_x, serialized_y, strict=True)],
                     initial=true_ary)

arraycontext/impl/pyopencl/taggable_cl_array.py

@@ -212,7 +212,7 @@ def zeros(
 
 def to_device(
             queue: cl.CommandQueue,
-            ary: np.ndarray[Any],
+            ary: np.ndarray[Any, Any],
             *, axes: tuple[Axis, ...] | None = None,
             tags: frozenset[Tag] = _EMPTY_TAG_SET,
             allocator: cla.Allocator | None = None,

arraycontext/impl/pytato/compile.py

@@ -100,12 +100,12 @@ def __hash__(self):
 
 @dataclass(frozen=True, eq=True)
 class ScalarInputDescriptor(AbstractInputDescriptor):
-    dtype: np.dtype
+    dtype: np.dtype[Any]
 
 
 @dataclass(frozen=True, eq=True)
 class LeafArrayDescriptor(AbstractInputDescriptor):
-    dtype: np.dtype
+    dtype: np.dtype[Any]
     shape: pt.array.ShapeType
 
 # }}}