[mlir][py] Enable AsmState overload for operation.

Author: jpienaar

mlir/lib/Bindings/Python/IRCore.cpp

@@ -3479,6 +3479,8 @@ void mlir::python::populateIRCore(py::module &m) {
 
   py::class_<PyAsmState>(m, "AsmState", py::module_local())
       .def(py::init<PyValue &, bool>(), py::arg("value"),
+           py::arg("use_local_scope") = false)
+      .def(py::init<PyOperationBase &, bool>(), py::arg("op"),
            py::arg("use_local_scope") = false);
 
   //----------------------------------------------------------------------------

mlir/lib/Bindings/Python/IRModule.h

@@ -759,6 +759,16 @@ class PyAsmState {
       mlirOpPrintingFlagsUseLocalScope(flags);
     state = mlirAsmStateCreateForValue(value, flags);
   }
+
+  PyAsmState(PyOperationBase &operation, bool useLocalScope) {
+    flags = mlirOpPrintingFlagsCreate();
+    // The OpPrintingFlags are not exposed Python side, create locally and
+    // associate lifetime with the state.
+    if (useLocalScope)
+      mlirOpPrintingFlagsUseLocalScope(flags);
+    state =
+        mlirAsmStateCreateForOperation(operation.getOperation().get(), flags);
+  }
   ~PyAsmState() {
     mlirOpPrintingFlagsDestroy(flags);
   }

mlir/test/python/ir/value.py

@@ -165,8 +165,8 @@ def testValuePrintAsOperand():
             # CHECK: Value(%[[VAL2:.*]] = "custom.op2"() : () -> i32)
             print(value2)
 
-            f = func.FuncOp("test", ([i32, i32], []))
-            entry_block1 = Block.create_at_start(f.operation.regions[0], [i32, i32])
+            topFn = func.FuncOp("test", ([i32, i32], []))
+            entry_block1 = Block.create_at_start(topFn.operation.regions[0], [i32, i32])
 
             with InsertionPoint(entry_block1):
                 value3 = Operation.create("custom.op3", results=[i32]).results[0]
@@ -201,7 +201,7 @@ def testValuePrintAsOperand():
 
         print("With AsmState")
         # CHECK-LABEL: With AsmState
-        state = AsmState(value3, use_local_scope=True)
+        state = AsmState(topFn.operation, use_local_scope=True)
         # CHECK: %0
         print(value3.get_name(state=state))
         # CHECK: %1