[Mlir-commits] [mlir] [mlir][python] Expose AsmState python side. (PR #66819)

[Maksim Levental]

================
@@ -3425,19 +3425,35 @@ void mlir::python::populateIRCore(py::module &m) {
           kValueDunderStrDocstring)
       .def(
           "get_name",
-          [](PyValue &self, bool useLocalScope) {
+          [](PyValue &self, std::optional<bool> useLocalScope,
+             std::optional<std::reference_wrapper<PyAsmState>> state) {
             PyPrintAccumulator printAccum;
-            MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
-            if (useLocalScope)
-              mlirOpPrintingFlagsUseLocalScope(flags);
-            MlirAsmState state = mlirAsmStateCreateForValue(self.get(), flags);
-            mlirValuePrintAsOperand(self.get(), state, printAccum.getCallback(),
+            MlirOpPrintingFlags flags;
+            MlirAsmState valueState;
+            // Use state if provided, else create a new state.
+            if (state) {
+              valueState = state.value().get().get();
+              // Don't allow setting using local scope and state at same time.
+              if (useLocalScope)
+                throw py::value_error(
+                    "setting AsmState and local scope together not supported");
+            } else {
+              flags = mlirOpPrintingFlagsCreate();
+              if (useLocalScope.value_or(false))
+                mlirOpPrintingFlagsUseLocalScope(flags);
+              valueState = mlirAsmStateCreateForValue(self.get(), flags);
+            }
+            mlirValuePrintAsOperand(self.get(), valueState, printAccum.getCallback(),
                                     printAccum.getUserData());
-            mlirOpPrintingFlagsDestroy(flags);
-            mlirAsmStateDestroy(state);
+            // Release state if allocated locally.
+            if (!state) {
+              mlirOpPrintingFlagsDestroy(flags);
+              mlirAsmStateDestroy(valueState);
+            }
             return printAccum.join();
           },
-          py::arg("use_local_scope") = false, kGetNameAsOperand)
+          py::arg("use_local_scope") = std::nullopt,
+          py::arg("state") = std::nullopt, kGetNameAsOperand)
----------------
makslevental wrote:

Ya I meant to `sed` it at some point but didn't get around to it

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