misc: use experimental TypeForm feature of Pyright

pyproject.toml

@@ -67,6 +67,7 @@ build-backend = "setuptools.build_meta"
 reportImportCycles = false
 reportMissingModuleSource = false
 enableTypeIgnoreComments = false
+enableExperimentalFeatures = true
 typeCheckingMode = "strict"
 "include" = ["docs", "xdsl", "tests", "bench"]
 "exclude" = [

tests/pattern_rewriter/test_pattern_rewriter.py

@@ -8,6 +8,7 @@
 from xdsl.dialects import test
 from xdsl.dialects.arith import Addi, Arith, Constant, Muli
 from xdsl.dialects.builtin import (
+    AnyIntegerAttr,
     Builtin,
     IndexType,
     IntegerAttr,
@@ -245,7 +246,7 @@ def test_recursive_rewriter():
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
         def match_and_rewrite(self, op: Constant, rewriter: PatternRewriter):
-            if not isa(op_val := op.value, IntegerAttr):
+            if not isa(op_val := op.value, AnyIntegerAttr):
                 return
             val = op_val.value.data
             if val == 0 or val == 1:
@@ -288,7 +289,7 @@ def test_recursive_rewriter_reversed():
     class Rewrite(RewritePattern):
         @op_type_rewrite_pattern
         def match_and_rewrite(self, op: Constant, rewriter: PatternRewriter):
-            if not isa(op_val := op.value, IntegerAttr):
+            if not isa(op_val := op.value, AnyIntegerAttr):
                 return
             val = op_val.value.data
             if val == 0 or val == 1:

tests/test_is_satisfying_hint.py

@@ -360,7 +360,7 @@ class MyParamAttr(Generic[_T], ParametrizedAttribute):
 def test_parametrized_attribute():
     attr = MyParamAttr[IntAttr]([IntAttr(0)])
 
-    assert isa(attr, MyParamAttr)
+    # `assert isa(attr, MyParamAttr)` not supported: use isinstance instead
     assert isa(attr, MyParamAttr[IntAttr])
     assert isa(attr, MyParamAttr[IntAttr | FloatData])
     assert not isa(attr, MyParamAttr[FloatData])

xdsl/transforms/csl_stencil_bufferize.py

@@ -125,7 +125,7 @@ def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter,
             zip(op.receive_chunk.block.args, buf_apply_op.receive_chunk.block.args)
         ):
             # arg0 has special meaning and does not need a `to_tensor` op
-            if isattr(old_arg.type, TensorType) and idx != 0:
+            if isattr(old_arg.type, AnyTensorTypeConstr) and idx != 0:
                 rewriter.insert_op(
                     # ensure iter_arg is writable
                     t := to_tensor_op(arg, writable=idx == 2),
@@ -139,7 +139,7 @@ def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter,
         for idx, (old_arg, arg) in enumerate(
             zip(op.done_exchange.block.args, buf_apply_op.done_exchange.block.args)
         ):
-            if isattr(old_arg.type, TensorType):
+            if isattr(old_arg.type, AnyTensorTypeConstr):
                 rewriter.insert_op(
                     # ensure iter_arg is writable
                     t := to_tensor_op(arg, writable=idx == 1),

xdsl/transforms/lower_csl_stencil.py

@@ -5,6 +5,7 @@
 from xdsl.dialects import arith, func, memref, stencil
 from xdsl.dialects.builtin import (
     AnyFloatAttr,
+    AnyMemRefTypeConstr,
     ArrayAttr,
     DenseIntOrFPElementsAttr,
     Float32Type,
@@ -440,7 +441,7 @@ def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter,
             return
 
         if (
-            not isattr(accumulator.type, memref.MemRefType)
+            not isattr(accumulator.type, AnyMemRefTypeConstr)
             or not isinstance(op.accumulator, OpResult)
             or not isinstance(alloc := op.accumulator.op, memref.Alloc)
         ):

xdsl/transforms/shape_inference_patterns/stencil.py

@@ -37,7 +37,7 @@ def update_result_size(
         apply = value.owner
         res_types = (cast(TempType[Attribute], r.type) for r in apply.res)
         newsize = reduce(
-            lambda l, r: l | r,
+            StencilBoundsAttr.union,
             (
                 size,
                 *(

xdsl/transforms/varith_transformations.py

@@ -14,7 +14,6 @@
     op_type_rewrite_pattern,
 )
 from xdsl.rewriter import InsertPoint
-from xdsl.utils.hints import isa
 
 # map the arith operation to the right varith op:
 ARITH_TO_VARITH_TYPE_MAP: dict[
@@ -138,14 +137,14 @@ def match_and_rewrite(self, op: varith.VarithOp, rewriter: PatternRewriter, /):
         # iterate over operands of the varith op:
         for inp in op.operands:
             # if the input happens to be the right arith op:
-            if isa(inp.owner, target_arith_type):
+            if isinstance(inp.owner, target_arith_type):
                 # fuse the operands of the arith op into the new varith op
                 new_operands.append(inp.owner.lhs)
                 new_operands.append(inp.owner.rhs)
                 # check if the old arith op can be erased
                 possibly_erased_ops.append(inp.owner)
             # if the parent op is a varith op of the same type as us
-            elif isa(inp.owner, type(op)):
+            elif isinstance(inp.owner, type(op)):
                 # include all operands of that
                 new_operands.extend(inp.owner.operands)
                 # check the old varith op for usages

xdsl/utils/hints.py

@@ -2,6 +2,7 @@
 from collections.abc import Iterable, Sequence
 from inspect import isclass
 from typing import (
+    TYPE_CHECKING,
     Annotated,
     Any,
     Generic,
@@ -17,10 +18,13 @@
 from xdsl.ir import ParametrizedAttribute
 from xdsl.utils.exceptions import VerifyException
 
+if TYPE_CHECKING:
+    from typing_extensions import TypeForm
+
 _T = TypeVar("_T")
 
 
-def isa(arg: Any, hint: type[_T]) -> TypeGuard[_T]:
+def isa(arg: Any, hint: "TypeForm[_T]") -> TypeGuard[_T]:
     from xdsl.irdl import ConstraintContext
 
     """

xdsl/utils/isattr.py

@@ -1,13 +1,16 @@
-from typing import Any, TypeGuard
+from typing import TYPE_CHECKING, Any, TypeGuard
 
 from xdsl.ir import AttributeInvT
 from xdsl.irdl import GenericAttrConstraint
 from xdsl.utils.exceptions import VerifyException
 from xdsl.utils.hints import isa
 
+if TYPE_CHECKING:
+    from typing_extensions import TypeForm
+
 
 def isattr(
-    arg: Any, hint: type[AttributeInvT] | GenericAttrConstraint[AttributeInvT]
+    arg: Any, hint: "TypeForm[AttributeInvT]" | GenericAttrConstraint[AttributeInvT]
 ) -> TypeGuard[AttributeInvT]:
     """
     A helper method to check whether a given attribute has a given type or conforms to a