1 files changed, 10 insertions, 0 deletions
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 23338f7..3ca7dd8 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/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);
   }