fix more pyright issues

majosm · majosm · commit e7f77b5e84a5 · 2025-06-04T15:00:56.000-05:00
diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json
@@ -22633,14 +22633,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportUnusedFunction",
-                "range": {
-                    "startColumn": 4,
-                    "endColumn": 8,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportMissingParameterType",
                 "range": {
@@ -24307,22 +24299,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 25,
-                    "endColumn": 29,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 25,
-                    "endColumn": 29,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportAttributeAccessIssue",
                 "range": {
@@ -24347,14 +24323,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 43,
-                    "endColumn": 47,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportArgumentType",
                 "range": {
@@ -24363,22 +24331,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 37,
-                    "endColumn": 41,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 42,
-                    "endColumn": 50,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportArgumentType",
                 "range": {
@@ -24387,22 +24339,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 38,
-                    "endColumn": 42,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 43,
-                    "endColumn": 51,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportMissingParameterType",
                 "range": {
@@ -24615,22 +24551,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportUnusedFunction",
-                "range": {
-                    "startColumn": 4,
-                    "endColumn": 28,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnusedFunction",
-                "range": {
-                    "startColumn": 4,
-                    "endColumn": 30,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportMissingParameterType",
                 "range": {
@@ -24663,14 +24583,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportUnusedFunction",
-                "range": {
-                    "startColumn": 4,
-                    "endColumn": 23,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportGeneralTypeIssues",
                 "range": {
@@ -24735,14 +24647,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportUnusedFunction",
-                "range": {
-                    "startColumn": 4,
-                    "endColumn": 26,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportMissingParameterType",
                 "range": {
diff --git a/arraycontext/container/traversal.py b/arraycontext/container/traversal.py
@@ -777,7 +777,7 @@ def unflatten(
         checks are skipped.
     """
     # NOTE: https://github.com/python/mypy/issues/7057
-    offset = 0
+    offset: int = 0
     common_dtype = None
 
     def _unflatten(template_subary: ArrayOrContainer) -> ArrayOrContainer:
@@ -791,7 +791,8 @@ def _unflatten(template_subary: ArrayOrContainer) -> ArrayOrContainer:
             # {{{ validate subary
 
             if (
-                    isinstance(template_subary_c.size, Integer)
+                    isinstance(offset, Integer)
+                    and isinstance(template_subary_c.size, Integer)
                     and isinstance(ary.size, Integer)
                     and (offset + template_subary_c.size) > ary.size):
                 raise ValueError("'template' and 'ary' sizes do not match: "
@@ -816,6 +817,12 @@ def _unflatten(template_subary: ArrayOrContainer) -> ArrayOrContainer:
 
             # {{{ reshape
 
+            if not isinstance(template_subary_c.size, Integer):
+                raise NotImplementedError(
+                    "unflatten is not implemented for arrays with array-valued "
+                    "size.") from None
+
+            # FIXME: Not sure how to make the slicing part work for Array-valued sizes
             flat_subary = ary[offset:offset + template_subary_c.size]
             try:
                 subary = actx.np.reshape(flat_subary,
diff --git a/arraycontext/context.py b/arraycontext/context.py
@@ -582,7 +582,7 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
     def outline(self,
                 f: Callable[..., Any],
                 *,
-                id: Hashable | None = None) -> Callable[..., Any]:
+                id: Hashable | None = None) -> Callable[..., Any]:  # pyright: ignore[reportUnusedParameter]
         """
         Returns a drop-in-replacement for *f*. The behavior of the returned
         callable is specific to the derived class.
diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py
@@ -236,7 +236,7 @@ def outline(self,
                 f: Callable[..., Any],
                 *,
                 id: Hashable | None = None,
-                tags: frozenset[Tag] = frozenset()
+                tags: frozenset[Tag] = frozenset()  # pyright: ignore[reportCallInDefaultInitializer]
                 ) -> Callable[..., Any]:
         from pytato.tags import FunctionIdentifier
 
@@ -962,6 +962,7 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
         from .compile import LazilyJAXCompilingFunctionCaller
         return LazilyJAXCompilingFunctionCaller(self, f)
 
+    @override
     def transform_dag(self, dag: pytato.DictOfNamedArrays
                       ) -> pytato.DictOfNamedArrays:
         import pytato as pt
diff --git a/arraycontext/impl/pytato/outline.py b/arraycontext/impl/pytato/outline.py
@@ -45,15 +45,14 @@
 from arraycontext.context import (
     Array,
     ArrayOrContainer,
-    ArrayOrContainerTc,
     ArrayT,
 )
 from arraycontext.impl.pytato import _BasePytatoArrayContext
 
 
 def _get_arg_id_to_arg(args: tuple[object, ...],
                        kwargs: Mapping[str, object]
-                       ) -> immutabledict[tuple[object, ...], object]:
+                       ) -> immutabledict[tuple[object, ...], pt.Array]:
     """
     Helper for :meth:`OulinedCall.__call__`. Extracts mappings from argument id
     to argument values. See
@@ -104,7 +103,7 @@ def _get_output_arg_id_str(arg_id: tuple[object, ...]) -> str:
 
 
 def _get_arg_id_to_placeholder(
-        arg_id_to_arg: Mapping[tuple[object, ...], object],
+        arg_id_to_arg: Mapping[tuple[object, ...], pt.Array],
         prefix: str | None = None) -> immutabledict[tuple[object, ...], pt.Placeholder]:
     """
     Helper for :meth:`OulinedCall.__call__`. Constructs a :class:`pytato.Placeholder`
@@ -122,25 +121,25 @@ def _get_arg_id_to_placeholder(
 
 def _call_with_placeholders(
         f: Callable[..., object],
-        args: tuple[object],
+        args: tuple[object, ...],
         kwargs: Mapping[str, object],
         arg_id_to_placeholder: Mapping[tuple[object, ...], pt.Placeholder]) -> object:
     """
     Construct placeholders analogous to *args* and *kwargs* and call *f*.
     """
     def get_placeholder_replacement(
-            arg: ArrayOrContainerTc | Scalar | None, key: tuple[object, ...]
-            ) -> ArrayOrContainerTc | Scalar | None:
+            arg: ArrayOrContainer | Scalar | None, key: tuple[object, ...]
+            ) -> ArrayOrContainer | Scalar | None:
         if arg is None:
             return None
         elif np.isscalar(arg):
             return cast(Scalar, arg)
         elif isinstance(arg, pt.Array):
-            return cast(ArrayOrContainerTc, arg_id_to_placeholder[key])
+            return arg_id_to_placeholder[key]
         elif is_array_container_type(arg.__class__):
-            def _rec_to_placeholder(keys: tuple[object, ...], ary: ArrayT) -> ArrayT:
-                result = get_placeholder_replacement(ary, key + keys)
-                return cast(ArrayT, result)
+            def _rec_to_placeholder(
+                    keys: tuple[object, ...], ary: Array) -> Array:
+                return cast("Array", get_placeholder_replacement(ary, key + keys))
 
             return rec_keyed_map_array_container(_rec_to_placeholder, arg)
         else:
@@ -176,7 +175,7 @@ def _unpack_container(key: tuple[object, ...], ary: ArrayT) -> ArrayT:
 
 def _pack_output(
         output_template: ArrayOrContainer,
-        unpacked_output: Array | immutabledict[str, Array]
+        unpacked_output: pt.Array | immutabledict[str, pt.Array]
         ) -> ArrayOrContainer:
     """
     Pack *unpacked_output* into array containers according to *output_template*.
@@ -187,9 +186,9 @@ def _pack_output(
     elif is_array_container_type(output_template.__class__):
         assert isinstance(unpacked_output, immutabledict)
 
-        def _pack_into_container(key: tuple[object, ...], ary: Array) -> Array:
+        def _pack_into_container(key: tuple[object, ...], ary: Array) -> Array:  # pyright: ignore[reportUnusedParameter]
             key_str = _get_output_arg_id_str(key)
-            return unpacked_output[key_str]
+            return unpacked_output[key_str]  # type: ignore[index]
 
         return rec_keyed_map_array_container(_pack_into_container, output_template)
     else:
@@ -262,9 +261,6 @@ def __call__(self, *args: object, **kwargs: object) -> ArrayOrContainer:
         call_site_output = func_def(**call_bindings)
 
         assert isinstance(call_site_output, pt.Array | immutabledict)
-        # FIXME: pt.Array is not an actx Array
-        return _pack_output(cast("Array | immutabledict[str, Array]", output),
-                            cast("Array | immutabledict[str, Array]", call_site_output))
-
+        return _pack_output(output, call_site_output)
 
 # vim: foldmethod=marker
diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
@@ -24,8 +24,10 @@
 """
 
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass
 from functools import partial
+from typing import TypeAlias
 
 import numpy as np
 import pytest
@@ -34,6 +36,7 @@
 from pytools.tag import Tag
 
 from arraycontext import (
+    ArrayContext,
     BcastUntilActxArray,
     EagerJAXArrayContext,
     NumpyArrayContext,
@@ -58,6 +61,9 @@
 logger = logging.getLogger(__name__)
 
 
+ArrayContextFactory: TypeAlias = Callable[[], ArrayContext]
+
+
 # {{{ array context fixture
 
 class _PyOpenCLArrayContextForTests(PyOpenCLArrayContext):
@@ -1166,15 +1172,16 @@ def my_rhs(scale, vel):
     np.testing.assert_allclose(result.v, 3.14*v_x)
 
 
-def test_actx_compile_with_outlined_function(actx_factory):
+def test_actx_compile_with_outlined_function(actx_factory: ArrayContextFactory):
     actx = actx_factory()
     rng = np.random.default_rng()
 
     @actx.outline
-    def outlined_scale_and_orthogonalize(alpha, vel):
+    def outlined_scale_and_orthogonalize(alpha: float, vel: Velocity2D) -> Velocity2D:
         return scale_and_orthogonalize(alpha, vel)
 
-    def multi_scale_and_orthogonalize(alpha, vel1, vel2):
+    def multi_scale_and_orthogonalize(
+            alpha: float, vel1: Velocity2D, vel2: Velocity2D) -> np.ndarray:
         return make_obj_array([
             outlined_scale_and_orthogonalize(alpha, vel1),
             outlined_scale_and_orthogonalize(alpha, vel2)])
@@ -1193,10 +1200,10 @@ def multi_scale_and_orthogonalize(alpha, vel1, vel2):
 
     result1 = actx.to_numpy(scaled_speed1)
     result2 = actx.to_numpy(scaled_speed2)
-    np.testing.assert_allclose(result1.u, -3.14*v1_y)
-    np.testing.assert_allclose(result1.v, 3.14*v1_x)
-    np.testing.assert_allclose(result2.u, -3.14*v2_y)
-    np.testing.assert_allclose(result2.v, 3.14*v2_x)
+    np.testing.assert_allclose(result1.u, -3.14*v1_y)  # pyright: ignore[reportAttributeAccessIssue]
+    np.testing.assert_allclose(result1.v, 3.14*v1_x)  # pyright: ignore[reportAttributeAccessIssue]
+    np.testing.assert_allclose(result2.u, -3.14*v2_y)  # pyright: ignore[reportAttributeAccessIssue]
+    np.testing.assert_allclose(result2.v, 3.14*v2_x)  # pyright: ignore[reportAttributeAccessIssue]
 
 # }}}
 
