[Mlir-commits] [mlir] [mlir][Python] downcast Value to BlockArgument or OpResult (PR #175264)

[Rahul Kayaith]

================
@@ -340,36 +340,47 @@ def testValueSetType():
 # CHECK-LABEL: TEST: testValueCasters
 @run
 def testValueCasters():
-    class NOPResult(OpResult):
-        def __init__(self, v):
-            super().__init__(v)
+    # check before registering casters
+    ctx = Context()
+    ctx.allow_unregistered_dialects = True
+    with Location.unknown(ctx):
+        i32 = IntegerType.get_signless(32)
+        module = Module.create()
+        with InsertionPoint(module.body):
+            value = Operation.create("custom.op0", results=[i32]).result
+            # CHECK: value OpResult(%0 = "custom.op0"() : () -> i32)
+            print("value", value)
 
-        def __str__(self):
-            return super().__str__().replace(Value.__name__, NOPResult.__name__)
+            @func.FuncOp.from_py_func(i32, i32)
+            def reduction(arg0, arg1):
+                # CHECK: arg0 BlockArgument(<block argument> of type 'i32' at index: 0)
+                print("arg0", arg0)
+                # CHECK: arg1 BlockArgument(<block argument> of type 'i32' at index: 1)
+                print("arg1", arg1)
 
-    class NOPValue(Value):
+    class NOPResult(OpResult):
         def __init__(self, v):
             super().__init__(v)
 
         def __str__(self):
-            return super().__str__().replace(Value.__name__, NOPValue.__name__)
+            return super().__str__().replace(OpResult.__name__, NOPResult.__name__)
 
     class NOPBlockArg(BlockArgument):
         def __init__(self, v):
             super().__init__(v)
 
         def __str__(self):
-            return super().__str__().replace(Value.__name__, NOPBlockArg.__name__)
+            return (
+                super().__str__().replace(BlockArgument.__name__, NOPBlockArg.__name__)
+            )
 
     @register_value_caster(IntegerType.static_typeid)
-    def cast_int(v) -> Value:
+    def cast_int(v) -> NOPResult | NOPBlockArg:
         print("in caster", v.__class__.__name__)
         if isinstance(v, OpResult):
             return NOPResult(v)
         if isinstance(v, BlockArgument):
             return NOPBlockArg(v)
-        elif isinstance(v, Value):
-            return NOPValue(v)
----------------
rkayaith wrote:

super nit: assert / raise an error here?

https://github.com/llvm/llvm-project/pull/175264