fix some pyright errors

Author: majosm

arraycontext/impl/pytato/__init__.py

@@ -58,6 +58,7 @@
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
+from typing_extensions import override
 
 from pytools import memoize_method
 from pytools.tag import Tag, ToTagSetConvertible, normalize_tags
@@ -230,6 +231,7 @@ def get_target(self):
 
     # }}}
 
+    @override
     def outline(self,
                 f: Callable[..., Any],
                 *,
@@ -582,8 +584,7 @@ def _to_frozen(key: tuple[Any, ...], ary) -> TaggableCLArray:
         pt_dict_of_named_arrays = pt.make_dict_of_named_arrays(
                 key_to_pt_arrays)
 
-        pt_dict_of_named_arrays = pt.transform.Deduplicator()(
-            pt_dict_of_named_arrays)
+        pt_dict_of_named_arrays = pt.deduplicate(pt_dict_of_named_arrays)
 
         # FIXME: Remove this if/when _normalize_pt_expr gets support for functions
         pt_dict_of_named_arrays = pt.tag_all_calls_to_be_inlined(

arraycontext/impl/pytato/compile.py

@@ -189,8 +189,7 @@ def _to_input_for_compiled(
     """
     from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array
     if isinstance(ary, pt.Array):
-        dag = pt.transform.Deduplicator()(
-            pt.make_dict_of_named_arrays({"_actx_out": ary}))
+        dag = pt.deduplicate(pt.make_dict_of_named_arrays({"_actx_out": ary}))
         # Transform the DAG to give metadata inference a chance to do its job
         return actx.transform_dag(dag)["_actx_out"].expr
     elif isinstance(ary, TaggableCLArray):

arraycontext/impl/pytato/outline.py

@@ -133,7 +133,9 @@ def get_placeholder_replacement(
             return arg_id_to_placeholder[key]
         elif is_array_container_type(arg.__class__):
             def _rec_to_placeholder(keys: tuple[Any, ...], ary: pt.Array) -> pt.Array:
-                return get_placeholder_replacement(ary, key + keys)
+                result = get_placeholder_replacement(ary, key + keys)
+                assert isinstance(result, pt.Array)
+                return result
 
             return rec_keyed_map_array_container(_rec_to_placeholder, arg)
         else:
@@ -206,9 +208,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> ArrayOrContainer:
 
             prefixed_output = _call_with_placeholders(
                 self.f, args, kwargs, arg_id_to_prefixed_placeholder)
-            unpacked_prefixed_output = pt.transform.Deduplicator()(
-                pt.make_dict_of_named_arrays(
-                    _unpack_output(prefixed_output)))
+            unpacked_prefixed_output = pt.deduplicate(
+                pt.make_dict_of_named_arrays(_unpack_output(prefixed_output)))
 
             prefixed_placeholders = frozenset(
                 arg_id_to_prefixed_placeholder.values())
@@ -225,9 +226,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> ArrayOrContainer:
         arg_id_to_placeholder = _get_arg_id_to_placeholder(arg_id_to_arg)
 
         output = _call_with_placeholders(self.f, args, kwargs, arg_id_to_placeholder)
-        unpacked_output = pt.transform.Deduplicator()(
-            pt.make_dict_of_named_arrays(
-                _unpack_output(output)))
+        unpacked_output = pt.deduplicate(
+            pt.make_dict_of_named_arrays(_unpack_output(output)))
         if len(unpacked_output) == 1 and "_" in unpacked_output:
             ret_type = pt.function.ReturnType.ARRAY
         else:

arraycontext/impl/pytato/utils.py

@@ -41,10 +41,11 @@
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, cast
 
+from typing_extensions import override
+
 import pytools
 from pytato.analysis import get_num_call_sites
 from pytato.array import (
-    AbstractResultWithNamedArrays,
     Array,
     Axis as PtAxis,
     DataWrapper,
@@ -58,8 +59,8 @@
 from pytato.transform import (
     ArrayOrNames,
     CopyMapper,
-    Deduplicator,
     TransformMapperCache,
+    deduplicate,
 )
 from pytools import UniqueNameGenerator, memoize_method
 
@@ -95,6 +96,7 @@ def __init__(
         self.vng = UniqueNameGenerator()
         self.seen_inputs: set[str] = set()
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         if expr.name is not None:
             if expr.name in self.seen_inputs:
@@ -114,13 +116,16 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array:
                     axes=expr.axes,
                     tags=expr.tags)
 
+    @override
     def map_size_param(self, expr: SizeParam) -> Array:
         raise NotImplementedError
 
+    @override
     def map_placeholder(self, expr: Placeholder) -> Array:
         raise ValueError("Placeholders cannot appear in"
                          " DatawrapperToBoundPlaceholderMapper.")
 
+    @override
     def map_function_definition(
             self, expr: FunctionDefinition) -> FunctionDefinition:
         raise ValueError("Function definitions cannot appear in"
@@ -131,7 +136,7 @@ def map_function_definition(
 # definitions can't contain non-argument placeholders
 def _normalize_pt_expr(
         expr: DictOfNamedArrays
-        ) -> tuple[Array | AbstractResultWithNamedArrays, Mapping[str, Any]]:
+        ) -> tuple[DictOfNamedArrays, Mapping[str, Any]]:
     """
     Returns ``(normalized_expr, bound_arguments)``.  *normalized_expr* is a
     normalized form of *expr*, with all instances of
@@ -141,7 +146,7 @@ def _normalize_pt_expr(
     Deterministic naming of placeholders permits more effective caching of
     equivalent graphs.
     """
-    expr = Deduplicator()(expr)
+    expr = deduplicate(expr)
 
     if get_num_call_sites(expr):
         raise NotImplementedError(
@@ -150,7 +155,7 @@ def _normalize_pt_expr(
 
     normalize_mapper = _DatawrapperToBoundPlaceholderMapper()
     normalized_expr = normalize_mapper(expr)
-    assert isinstance(normalized_expr, AbstractResultWithNamedArrays)
+    assert isinstance(normalized_expr, DictOfNamedArrays)
     return normalized_expr, normalize_mapper.bound_arguments
 
 
@@ -188,6 +193,7 @@ def __init__(self, actx: ArrayContext) -> None:
         super().__init__()
         self.actx = actx
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         import numpy as np
 
@@ -220,6 +226,7 @@ def __init__(self, actx: ArrayContext) -> None:
         super().__init__()
         self.actx = actx
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         import numpy as np