more pyright fixes

majosm · majosm · commit 74675c7b02b7 · 2025-06-16T15:27:04.000-05:00
diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json
@@ -11653,14 +11653,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 53,
-                    "endColumn": 83,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportPrivateUsage",
                 "range": {
@@ -11765,30 +11757,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportArgumentType",
-                "range": {
-                    "startColumn": 21,
-                    "endColumn": 60,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportArgumentType",
-                "range": {
-                    "startColumn": 49,
-                    "endColumn": 67,
-                    "lineCount": 2
-                }
-            },
-            {
-                "code": "reportArgumentType",
-                "range": {
-                    "startColumn": 53,
-                    "endColumn": 68,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnknownMemberType",
                 "range": {
@@ -11829,14 +11797,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportArgumentType",
-                "range": {
-                    "startColumn": 16,
-                    "endColumn": 58,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportPrivateUsage",
                 "range": {
@@ -11885,22 +11845,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportArgumentType",
-                "range": {
-                    "startColumn": 12,
-                    "endColumn": 51,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportArgumentType",
-                "range": {
-                    "startColumn": 45,
-                    "endColumn": 63,
-                    "lineCount": 2
-                }
-            },
             {
                 "code": "reportUnknownMemberType",
                 "range": {
@@ -13045,14 +12989,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 53,
-                    "endColumn": 83,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportArgumentType",
                 "range": {
@@ -17283,14 +17219,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportImplicitOverride",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 24,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportPrivateUsage",
                 "range": {
@@ -17299,22 +17227,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportImplicitOverride",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 22,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportImplicitOverride",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 23,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnannotatedClassAttribute",
                 "range": {
@@ -17331,14 +17243,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportImplicitOverride",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 24,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnannotatedClassAttribute",
                 "range": {
@@ -17347,14 +17251,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportImplicitOverride",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 24,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportPrivateUsage",
                 "range": {
diff --git a/arraycontext/context.py b/arraycontext/context.py
@@ -587,7 +587,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
@@ -238,7 +238,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
 
@@ -976,6 +976,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,7 +186,7 @@ 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]
 
@@ -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]
 
 # }}}
 
