Repository: inducer/pytato (language: python)

Issue:
Typing of mappers
> I thought about this a little more, and it got me wondering: does it actually make sense to require `CopyMapper`'s `rec`/`__call__` to return the same type of thing as it takes in? In other words, couldn't a mapper in theory replace an `Array` with a `DictOfNamedArrays`, and vice versa?
>
> Edit: Hmm, maybe not. I was thinking of function inlining as an example, but I guess that's not quite what happens there.

_Originally posted by @majosm in https://github.com/inducer/pytato/issues/516#issuecomment-2225831659_
            


Discussion / hints:
You're right, this could happen. The only clean-ish way I can think of is to have two different mapper types: one that promises to not change types, and another that allows it? But that's a lot of ceremony just to please a type checker.
@kaushikcfd any opinions here?
I'm having trouble thinking of an example of this happening now, after my original thought about the function inlining evaporated. Since the interface for operating on an `Array` is different from that of an `AbstractResultWithNamedArrays`, it seems unlikely that you'd directly swap them. Realistically you'd only ever be replacing a `NamedArray` with some other kind of `Array`, which is allowed. Do you buy that, or is there a case that I'm not thinking of?
Typical sub-classes of CopyMapper are expression-rewriters for which returning DictOfNamedArrays from an Array is uncommon. So I think the current typing is sufficiently sound for practical reasons.
Hmm. Based on a question from @majosm this morning, I fired off #516. Now that I've actually read [PEP484](https://peps.python.org/pep-0484/#generics)  in more detail, this may actually not be what we want, as that makes the typing more demanding than before.

FYI @majosm the relevant example from PEP 484 is this:

> ```
> AnyStr = TypeVar('AnyStr', Text, bytes)
>
> def concat(x: AnyStr, y: AnyStr) -> AnyStr:
>     return x + y
>
> class MyStr(str): ...
>
> x = concat(MyStr('apple'), MyStr('pie'))
> ```
> [...]
> The call is valid but the type variable `AnyStr` will be set to `str` and not `MyStr`. In effect, the inferred type of the return value assigned to x will also be `str`.

So the net effect of the annotation before the (misguided, I think) changes in #516 was exactly what @kaushikcfd described: E.g. you feed in an `IndexLambda`, and you're promised that an `Array` comes out (but not a `DictOfNamedArrays` or an `IndexLambda`). Meanwhile, post #516, the annotation promises that an `IndexLambda` will come out, which is a promise that may be broken by term rewriting.

tl;dr: I'll be reverting #516.

Fix patch:
diff --git a/pytato/transform/__init__.py b/pytato/transform/__init__.py
index 3031f98e7..76a9ab14b 100644
--- a/pytato/transform/__init__.py
+++ b/pytato/transform/__init__.py
@@ -54,9 +54,11 @@
 
 
 ArrayOrNames = Union[Array, AbstractResultWithNamedArrays]
-MappedT = TypeVar("MappedT", bound=ArrayOrNames)
+MappedT = TypeVar("MappedT",
+                  Array, AbstractResultWithNamedArrays, ArrayOrNames)
 CombineT = TypeVar("CombineT")  # used in CombineMapper
-CopyMapperResultT = TypeVar("CopyMapperResultT", bound=ArrayOrNames)
+CopyMapperResultT = TypeVar("CopyMapperResultT",  # used in CopyMapper
+                            Array, AbstractResultWithNamedArrays, ArrayOrNames)
 CachedMapperT = TypeVar("CachedMapperT")  # used in CachedMapper
 IndexOrShapeExpr = TypeVar("IndexOrShapeExpr")
 R = FrozenSet[Array]


Test patch:
diff --git a/test/testlib.py b/test/testlib.py
index f5e19abb4..73daf101f 100644
--- a/test/testlib.py
+++ b/test/testlib.py
@@ -259,6 +259,7 @@ def get_random_pt_dag(seed: int,
         additional_generators = []
 
     from testlib import RandomDAGContext, make_random_dag
+    from typing import cast
 
     rdagc_comm = RandomDAGContext(np.random.default_rng(seed=seed),
             axis_len=axis_len, use_numpy=False,
@@ -277,7 +278,8 @@ def make_dws_placeholder(expr: pt.transform.ArrayOrNames
             else:
                 return expr
 
-        dag = pt.transform.map_and_copy(dag, make_dws_placeholder)
+        dag = cast(pt.DictOfNamedArrays,
+                   pt.transform.map_and_copy(dag, make_dws_placeholder))
 
     return dag