Mapper: allow 'expr' to be subclasses of ArrayOrNames

Author: kaushikcfd

pytato/analysis/__init__.py

@@ -76,9 +76,7 @@ def rec(self, expr: ArrayOrNames) -> None:  # type: ignore
         if id(expr) in self._visited_ids:
             return
 
-        # type-ignore reason: super().rec expects either 'Array' or
-        # 'AbstractResultWithNamedArrays', passed 'ArrayOrNames'
-        super().rec(expr)  # type: ignore
+        super().rec(expr)
         self._visited_ids.add(id(expr))
 
     def map_index_lambda(self, expr: IndexLambda) -> None:

pytato/distributed.py

@@ -578,8 +578,7 @@ def rec(self, expr: ArrayOrNames) -> FrozenSet[Array]:  # type: ignore[override]
         try:
             return self.cache[expr]
         except KeyError:
-            # type-ignore reason: type not compatible with super.rec() type
-            result: FrozenSet[Array] = super().rec(expr)  # type: ignore[type-var]
+            result: FrozenSet[Array] = super().rec(expr)
             self.cache[expr] = result
             return result
 

pytato/partition.py

@@ -72,7 +72,7 @@
 
 
 ArrayOrNames = Union[Array, AbstractResultWithNamedArrays]
-T = TypeVar("T", Array, AbstractResultWithNamedArrays)
+T = TypeVar("T", bound=ArrayOrNames)
 PartId = Hashable
 
 

pytato/target/python/numpy_like.py

@@ -549,8 +549,7 @@ def generate_numpy_like(expr: Union[Array, Mapping[str, Array], DictOfNamedArray
         numpy_backend=target.numpy_like_module_name_shorthand,
         numpy="np",
         vng=var_name_gen)
-    # type-ignore-reason: https://github.com/inducer/pytato/issues/236
-    result_var = cgen_mapper(expr)  # type: ignore[type-var]
+    result_var = cgen_mapper(expr)
 
     lines = cgen_mapper.lines
     lines.append(ast.Return(ast.Name(result_var)))

pytato/transform/__init__.py

@@ -50,11 +50,11 @@
 if TYPE_CHECKING:
     from pytato.distributed import DistributedSendRefHolder, DistributedRecv
 
-T = TypeVar("T", Array, AbstractResultWithNamedArrays)
+ArrayOrNames = Union[Array, AbstractResultWithNamedArrays]
+T = TypeVar("T", bound=ArrayOrNames)
 CombineT = TypeVar("CombineT")  # used in CombineMapper
 CachedMapperT = TypeVar("CachedMapperT")  # used in CachedMapper
 IndexOrShapeExpr = TypeVar("IndexOrShapeExpr")
-ArrayOrNames = Union[Array, AbstractResultWithNamedArrays]
 R = FrozenSet[Array]
 
 __doc__ = """
@@ -378,7 +378,7 @@ def rec(self,
         except KeyError:
             result = Mapper.rec(self, expr,
                                 *args,
-                                **kwargs)  # type: ignore[type-var]
+                                **kwargs)
             self._cache[key] = result
             return result
 
@@ -554,8 +554,7 @@ def rec_idx_or_size_tuple(self, situp: Tuple[IndexOrShapeExpr, ...]
     def rec(self, expr: ArrayOrNames) -> CombineT:  # type: ignore
         if expr in self.cache:
             return self.cache[expr]
-        # type-ignore reason: type not compatible with super.rec() type
-        result: CombineT = super().rec(expr)  # type: ignore
+        result: CombineT = super().rec(expr)
         self.cache[expr] = result
         return result
 
@@ -961,9 +960,7 @@ def rec(self, expr: ArrayOrNames) -> None:  # type: ignore
         if id(expr) in self._visited_ids:
             return
 
-        # type-ignore reason: super().rec expects either 'Array' or
-        # 'AbstractResultWithNamedArrays', passed 'ArrayOrNames'
-        super().rec(expr)  # type: ignore
+        super().rec(expr)
         self._visited_ids.add(id(expr))
 
 # }}}
@@ -1061,8 +1058,7 @@ def __init__(self, nsuccessors: Mapping[Array, int]):
     def rec(self, expr: ArrayOrNames) -> MPMSMaterializerAccumulator:  # type: ignore
         if expr in self.cache:
             return self.cache[expr]
-        # type-ignore reason: type not compatible with super.rec() type
-        result: MPMSMaterializerAccumulator = super().rec(expr)  # type: ignore
+        result: MPMSMaterializerAccumulator = super().rec(expr)
         self.cache[expr] = result
         return result