Revert "Improve mapper type annotations"

This reverts commit dd6a35b.
inducer · Jul 12, 2024 · b528e25 · b528e25
1 parent dd6a35b
commit b528e25
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git 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]

diff --git 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