[mlir][python] Capture error diagnostics in exceptions

[rkayaith]

Summary

This updates most (all?) error-diagnostic-emitting python APIs to
capture error diagnostics and include them in the raised exception's
message:

>>> Operation.parse('"arith.addi"() : () -> ()'))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
mlir._mlir_libs.MLIRError: Unable to parse operation assembly:
error: "-":1:1: 'arith.addi' op requires one result
 note: "-":1:1: see current operation: "arith.addi"() : () -> ()

The diagnostic information is available on the exception for users who
may want to customize the error message:

>>> try:
...   Operation.parse('"arith.addi"() : () -> ()')
... except MLIRError as e:
...   print(e.message)
...   print(e.error_diagnostics)
...   print(e.error_diagnostics[0].message)
... 
Unable to parse operation assembly
[<mlir._mlir_libs._mlir.ir.DiagnosticInfo object at 0x7fed32bd6b70>]
'arith.addi' op requires one result

Error diagnostics captured in exceptions aren't propagated to diagnostic
handlers, to avoid double-reporting of errors. The context-level
emit_error_diagnostics option can be used to revert to the old
behaviour, causing error diagnostics to be reported to handlers instead
of as part of exceptions.

API changes:

• Operation.verify now raises an exception on verification failure, instead of returning false
• The exception raised by the following methods has been changed to MLIRError:
  • PassManager.run
  • {Module,Operation,Type,Attribute}.parse
  • {RankedTensorType,UnrankedTensorType}.get
  • {MemRefType,UnrankedMemRefType}.get
  • VectorType.get
  • FloatAttr.get

closes #60595

depends on D144804, D143830

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

3380–3381

I couldn't find a way to define an exception with custom __init__ through pybind, so it's defined in python + imported here. Another option would be to use the C API.

[rkayaith]

update

[rkayaith]

missed some methods

[rkayaith]

Since I only changed error-emitting APIs to raise MLIRError, the exception type used is somewhat inconsistent, e.g. ComplexType.get still raises ValueError. Not sure what to do about those (raise MLIRError everywhere even if there's no diagnostics?)

[stellaraccident]

Nice, thanks! Fine to land as a stretch improvement but I think we need to further sugar the new exception class to avoid exposing implementation details.

mlir/python/mlir/_mlir_libs/__init__.py

109

I think we may want to further typify this vs having deeply nested tuples. Can be fine in a follow-up to sugar. Would make string printing easier and would leak fewer details to callers who want to extract the diagnostics.

[rkayaith]

Replace diagnostic tuples with a DiagnosticInfo class, which just mirrors the
main Diagnostic class but is safe to access outside handlers.

mlir/python/mlir/_mlir_libs/__init__.py

109

Took a stab at this, let me know if this is what you had in mind.

[rkayaith]

rebase

mlir/lib/Bindings/Python/IRAttributes.cpp

public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAFloat;
   static constexpr const char *pyClassName = "FloatAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &type, double value, DefaultingPyLocation loc) {
+          PyMlirContext::ErrorCapture errors(loc->getContext());
           MlirAttribute attr = mlirFloatAttrDoubleGetChecked(loc, type, value);
-          // TODO: Rework error reporting once diagnostic engine is exposed
-          // in C API.
-          if (mlirAttributeIsNull(attr)) {
-            throw SetPyError(PyExc_ValueError,
-                             Twine("invalid '") +
-                                 py::repr(py::cast(type)).cast<std::string>() +
-                                 "' and expected floating point type.");
-          }
+          if (mlirAttributeIsNull(attr))
+            throw MLIRError("Invalid attribute", errors.take());
           return PyFloatAttribute(type.getContext(), attr);
         },
         py::arg("type"), py::arg("value"), py::arg("loc") = py::none(),
         "Gets an uniqued float point attribute associated to a type");
     c.def_static(
         "get_f32",
         [](double value, DefaultingPyMlirContext context) {
           MlirAttribute attr = mlirFloatAttrDoubleGet(

mlir/lib/Bindings/Python/IRCore.cpp

 
 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Debug.h"
+#include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 //#include "mlir-c/Registration.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <optional>
 #include <utility>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::StringRef;
 using llvm::Twine;
 
 //------------------------------------------------------------------------------
 // Docstrings (trivial, non-duplicated docstrings are included inline).
 //------------------------------------------------------------------------------
 
 static const char kContextParseTypeDocstring[] =
     R"(Parses the assembly form of a type.
 
-Returns a Type object or raises a ValueError if the type cannot be parsed.
+Returns a Type object or raises an MLIRError if the type cannot be parsed.
 
 See also: https://mlir.llvm.org/docs/LangRef/#type-system
 )";
 
 static const char kContextGetCallSiteLocationDocstring[] =
     R"(Gets a Location representing a caller and callsite)";
 
 static const char kContextGetFileLocationDocstring[] =
     R"(Gets a Location representing a file, line and column)";
 
 static const char kContextGetFusedLocationDocstring[] =
     R"(Gets a Location representing a fused location with optional metadata)";
 
 static const char kContextGetNameLocationDocString[] =
     R"(Gets a Location representing a named location with optional child location)";
 
 static const char kModuleParseDocstring[] =
     R"(Parses a module's assembly format from a string.
 
-Returns a new MlirModule or raises a ValueError if the parsing fails.
+Returns a new MlirModule or raises an MLIRError if the parsing fails.
 
 See also: https://mlir.llvm.org/docs/LangRef/
 )";
 
 static const char kOperationCreateDocstring[] =
     R"(Creates a new operation.
 
 Args:
auto deleteCallback = +[](void *userData) {
     pyHandlerObject.dec_ref();
   };
 
   pyHandler->registeredID = mlirContextAttachDiagnosticHandler(
       get(), handlerCallback, static_cast<void *>(pyHandler), deleteCallback);
   return pyHandlerObject;
 }
 
+MlirLogicalResult PyMlirContext::ErrorCapture::handler(MlirDiagnostic diag,
+                                                       void *userData) {
+  auto *self = static_cast<ErrorCapture *>(userData);
+  // Check if the context requested we emit errors instead of capturing them.
+  if (self->ctx->emitErrorDiagnostics)
+    return mlirLogicalResultFailure();
+
+  if (mlirDiagnosticGetSeverity(diag) != MlirDiagnosticError)
+    return mlirLogicalResultFailure();
+
+  self->errors.emplace_back(PyDiagnostic(diag).getInfo());
+  return mlirLogicalResultSuccess();
+}
+
 PyMlirContext &DefaultingPyMlirContext::resolve() {
   PyMlirContext *context = PyThreadContextEntry::getDefaultContext();
   if (!context) {
     throw SetPyError(
         PyExc_RuntimeError,
         "An MLIR function requires a Context but none was provided in the call "
         "or from the surrounding environment. Either pass to the function with "
         "a 'context=' argument or establish a default using 'with Context():'");
py::tuple PyDiagnostic::getNotes() {
   materializedNotes = py::tuple(numNotes);
   for (intptr_t i = 0; i < numNotes; ++i) {
     MlirDiagnostic noteDiag = mlirDiagnosticGetNote(diagnostic, i);
     (*materializedNotes)[i] = PyDiagnostic(noteDiag);
   }
   return *materializedNotes;
 }
 
+PyDiagnostic::DiagnosticInfo PyDiagnostic::getInfo() {
+  std::vector<DiagnosticInfo> notes;
+  for (py::handle n : getNotes())
+    notes.emplace_back(n.cast<PyDiagnostic>().getInfo());
+  return {getSeverity(), getLocation(), getMessage(), std::move(notes)};
+}
+
 //------------------------------------------------------------------------------
 // PyDialect, PyDialectDescriptor, PyDialects, PyDialectRegistry
 //------------------------------------------------------------------------------
 
 MlirDialect PyDialects::getDialectForKey(const std::string &key,
                                          bool attrError) {
   MlirDialect dialect = mlirContextGetOrLoadDialect(getContext()->get(),
                                                     {key.data(), key.size()});
PyOperationRef created = createInstance(std::move(contextRef), operation,
                                           std::move(parentKeepAlive));
   created->attached = false;
   return created;
 }
 
 PyOperationRef PyOperation::parse(PyMlirContextRef contextRef,
                                   const std::string &sourceStr,
                                   const std::string &sourceName) {
+  PyMlirContext::ErrorCapture errors(contextRef);
   MlirOperation op =
       mlirOperationCreateParse(contextRef->get(), toMlirStringRef(sourceStr),
                                toMlirStringRef(sourceName));
-  // TODO: Include error diagnostic messages in the exception message
   if (mlirOperationIsNull(op))
-    throw py::value_error(
-        "Unable to parse operation assembly (see diagnostics)");
+    throw MLIRError("Unable to parse operation assembly", errors.take());
   return PyOperation::createDetached(std::move(contextRef), op);
 }
 
 void PyOperation::checkValid() const {
   if (!valid) {
     throw SetPyError(PyExc_RuntimeError, "the operation has been invalidated");
   }
 }
void PyOperationBase::moveBefore(PyOperationBase &other) {
   PyOperation &operation = getOperation();
   PyOperation &otherOp = other.getOperation();
   operation.checkValid();
   otherOp.checkValid();
   mlirOperationMoveBefore(operation, otherOp);
   operation.parentKeepAlive = otherOp.parentKeepAlive;
 }
 
+bool PyOperationBase::verify() {
+  PyOperation &op = getOperation();
+  PyMlirContext::ErrorCapture errors(op.getContext());
+  if (!mlirOperationVerify(op.get()))
+    throw MLIRError("Verification failed", errors.take());
+  return true;
+}
+
 std::optional<PyOperationRef> PyOperation::getParentOperation() {
   checkValid();
   if (!isAttached())
     throw SetPyError(PyExc_ValueError, "Detached operations have no parent");
   MlirOperation operation = mlirOperationGetParentOperation(get());
   if (mlirOperationIsNull(operation))
     return {};
   return PyOperation::forOperation(getContext(), operation);
py::class_<PyDiagnostic>(m, "Diagnostic", py::module_local())
       .def_property_readonly("message", &PyDiagnostic::getMessage)
       .def_property_readonly("notes", &PyDiagnostic::getNotes)
       .def("__str__", [](PyDiagnostic &self) -> py::str {
         if (!self.isValid())
           return "<Invalid Diagnostic>";
         return self.getMessage();
       });
 
+  py::class_<PyDiagnostic::DiagnosticInfo>(m, "DiagnosticInfo",
+                                           py::module_local())
+      .def(py::init<>([](PyDiagnostic diag) { return diag.getInfo(); }))
+      .def_readonly("severity", &PyDiagnostic::DiagnosticInfo::severity)
+      .def_readonly("location", &PyDiagnostic::DiagnosticInfo::location)
+      .def_readonly("message", &PyDiagnostic::DiagnosticInfo::message)
+      .def_readonly("notes", &PyDiagnostic::DiagnosticInfo::notes)
+      .def("__str__",
+           [](PyDiagnostic::DiagnosticInfo &self) { return self.message; });
+
   py::class_<PyDiagnosticHandler>(m, "DiagnosticHandler", py::module_local())
       .def("detach", &PyDiagnosticHandler::detach)
       .def_property_readonly("attached", &PyDiagnosticHandler::isAttached)
       .def_property_readonly("had_error", &PyDiagnosticHandler::getHadError)
       .def("__enter__", &PyDiagnosticHandler::contextEnter)
       .def("__exit__", &PyDiagnosticHandler::contextExit);
 
   //----------------------------------------------------------------------------
py::class_<PyMlirContext>(m, "_BaseContext", py::module_local())
           },
           py::arg("operation_name"))
       .def(
           "append_dialect_registry",
           [](PyMlirContext &self, PyDialectRegistry &registry) {
             mlirContextAppendDialectRegistry(self.get(), registry);
           },
           py::arg("registry"))
+      .def_property("emit_error_diagnostics", nullptr,
+                    &PyMlirContext::setEmitErrorDiagnostics,
+                    "Emit error diagnostics to diagnostic handlers. By default "
+                    "error diagnostics are captured and reported through "
+                    "MLIRError exceptions.")
       .def("load_all_available_dialects", [](PyMlirContext &self) {
         mlirContextLoadAllAvailableDialects(self.get());
       });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialectDescriptor
   //----------------------------------------------------------------------------
   py::class_<PyDialectDescriptor>(m, "DialectDescriptor", py::module_local())
void mlir::python::populateIRCore(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of Module
   //----------------------------------------------------------------------------
   py::class_<PyModule>(m, "Module", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyModule::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyModule::createFromCapsule)
       .def_static(
           "parse",
-          [](const std::string moduleAsm, DefaultingPyMlirContext context) {
+          [](const std::string &moduleAsm, DefaultingPyMlirContext context) {
+            PyMlirContext::ErrorCapture errors(context->getRef());
             MlirModule module = mlirModuleCreateParse(
                 context->get(), toMlirStringRef(moduleAsm));
-            // TODO: Rework error reporting once diagnostic engine is exposed
-            // in C API.
-            if (mlirModuleIsNull(module)) {
-              throw SetPyError(
-                  PyExc_ValueError,
-                  "Unable to parse module assembly (see diagnostics)");
-            }
+            if (mlirModuleIsNull(module))
+              throw MLIRError("Unable to parse module assembly", errors.take());
             return PyModule::forModule(module).releaseObject();
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kModuleParseDocstring)
       .def_static(
           "create",
           [](DefaultingPyLocation loc) {
             MlirModule module = mlirModuleCreateEmpty(loc);
py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
            // Careful: Lots of arguments must match up with get_asm method.
            py::arg("binary") = false,
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
            py::arg("pretty_debug_info") = false,
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false,
            py::arg("assume_verified") = false, kOperationGetAsmDocstring)
-      .def(
-          "verify",
-          [](PyOperationBase &self) {
-            return mlirOperationVerify(self.getOperation());
-          },
-          "Verify the operation and return true if it passes, false if it "
-          "fails.")
+      .def("verify", &PyOperationBase::verify,
+           "Verify the operation. Raises MLIRError if verification fails, and "
+           "returns true otherwise.")
       .def("move_after", &PyOperationBase::moveAfter, py::arg("other"),
            "Puts self immediately after the other operation in its parent "
            "block.")
       .def("move_before", &PyOperationBase::moveBefore, py::arg("other"),
            "Puts self immediately before the other operation in its parent "
            "block.")
       .def(
           "detach_from_parent",
opViewClass.attr("parse") = classmethod(
 
         // Check if the expected operation was parsed, and cast to to the
         // appropriate `OpView` subclass if successful.
         // NOTE: This accesses attributes that have been automatically added to
         // `OpView` subclasses, and is not intended to be used on `OpView`
         // directly.
         std::string clsOpName =
             py::cast<std::string>(cls.attr("OPERATION_NAME"));
-        MlirStringRef parsedOpName =
+        MlirStringRef identifier =
             mlirIdentifierStr(mlirOperationGetName(*parsed.get()));
-        if (!mlirStringRefEqual(parsedOpName, toMlirStringRef(clsOpName)))
-          throw py::value_error(
-              "Expected a '" + clsOpName + "' op, got: '" +
-              std::string(parsedOpName.data, parsedOpName.length) + "'");
+        std::string_view parsedOpName(identifier.data, identifier.length);
+        if (clsOpName != parsedOpName)
+          throw MLIRError(Twine("Expected a '") + clsOpName + "' op, got: '" +
+                          parsedOpName + "'");
         return PyOpView::constructDerived(cls, *parsed.get());
       },
       py::arg("cls"), py::arg("source"), py::kw_only(),
       py::arg("source_name") = "", py::arg("context") = py::none(),
       "Parses a specific, generated OpView based on class level attributes");
 
   //----------------------------------------------------------------------------
   // Mapping of PyRegion.
py::class_<PyAttribute>(m, "Attribute", py::module_local())
       .def(py::init<PyAttribute &>(), py::arg("cast_from_type"),
            "Casts the passed attribute to the generic Attribute")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAttribute::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAttribute::createFromCapsule)
       .def_static(
           "parse",
           [](std::string attrSpec, DefaultingPyMlirContext context) {
+            PyMlirContext::ErrorCapture errors(context->getRef());
             MlirAttribute type = mlirAttributeParseGet(
                 context->get(), toMlirStringRef(attrSpec));
-            // TODO: Rework error reporting once diagnostic engine is exposed
-            // in C API.
-            if (mlirAttributeIsNull(type)) {
-              throw SetPyError(PyExc_ValueError,
-                               Twine("Unable to parse attribute: '") +
-                                   attrSpec + "'");
-            }
+            if (mlirAttributeIsNull(type))
+              throw MLIRError("Unable to parse attribute", errors.take());
             return PyAttribute(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
-          "Parses an attribute from an assembly form")
+          "Parses an attribute from an assembly form. Raises an MLIRError on "
+          "failure.")
       .def_property_readonly(
           "context",
           [](PyAttribute &self) { return self.getContext().getObject(); },
           "Context that owns the Attribute")
       .def_property_readonly("type",
                              [](PyAttribute &self) {
                                return PyType(self.getContext()->getRef(),
                                              mlirAttributeGetType(self));
py::class_<PyType>(m, "Type", py::module_local())
       // extend the backing context which owns the MlirType.
       .def(py::init<PyType &>(), py::arg("cast_from_type"),
            "Casts the passed type to the generic Type")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyType::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyType::createFromCapsule)
       .def_static(
           "parse",
           [](std::string typeSpec, DefaultingPyMlirContext context) {
+            PyMlirContext::ErrorCapture errors(context->getRef());
             MlirType type =
                 mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec));
-            // TODO: Rework error reporting once diagnostic engine is exposed
-            // in C API.
-            if (mlirTypeIsNull(type)) {
-              throw SetPyError(PyExc_ValueError,
-                               Twine("Unable to parse type: '") + typeSpec +
-                                   "'");
-            }
+            if (mlirTypeIsNull(type))
+              throw MLIRError("Unable to parse type", errors.take());
             return PyType(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kContextParseTypeDocstring)
       .def_property_readonly(
           "context", [](PyType &self) { return self.getContext().getObject(); },
           "Context that owns the Type")
       .def("__eq__", [](PyType &self, PyType &other) { return self == other; })

   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
   PyGlobalDebugFlag::bind(m);
 
   // Attribute builder getter.
   PyAttrBuilderMap::bind(m);
+
+  py::register_local_exception_translator([](std::exception_ptr p) {
+    // We can't define exceptions with custom fields through pybind, so instead
+    // the exception class is defined in python and imported here.
+    try {
+      if (p)
+        std::rethrow_exception(p);
+    } catch (const MLIRError &e) {
+      py::object obj = py::module_::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+                           .attr("MLIRError")(e.message, e.errorDiagnostics);
+      PyErr_SetObject(PyExc_Exception, obj.ptr());
+    }
+  });
 }

mlir/lib/Bindings/Python/IRModule.h

public:
   void contextExit(const pybind11::object &excType,
                    const pybind11::object &excVal,
                    const pybind11::object &excTb);
 
   /// Attaches a Python callback as a diagnostic handler, returning a
   /// registration object (internally a PyDiagnosticHandler).
   pybind11::object attachDiagnosticHandler(pybind11::object callback);
 
+  /// Controls whether error diagnostics should be propagated to diagnostic
+  /// handlers, instead of being captured by `ErrorCapture`.
+  void setEmitErrorDiagnostics(bool value) { emitErrorDiagnostics = value; }
+  struct ErrorCapture;
+
 private:
   PyMlirContext(MlirContext context);
   // Interns the mapping of live MlirContext::ptr to PyMlirContext instances,
   // preserving the relationship that an MlirContext maps to a single
   // PyMlirContext wrapper. This could be replaced in the future with an
   // extension mechanism on the MlirContext for stashing user pointers.
   // Note that this holds a handle, which does not imply ownership.
   // Mappings will be removed when the context is destructed.
private:
   // Interns all live operations associated with this context. Operations
   // tracked in this map are valid. When an operation is invalidated, it is
   // removed from this map, and while it still exists as an instance, any
   // attempt to access it will raise an error.
   using LiveOperationMap =
       llvm::DenseMap<void *, std::pair<pybind11::handle, PyOperation *>>;
   LiveOperationMap liveOperations;
 
+  bool emitErrorDiagnostics = false;
+
   MlirContext context;
   friend class PyModule;
   friend class PyOperation;
 };
 
 /// Used in function arguments when None should resolve to the current context
 /// manager set instance.
 class DefaultingPyMlirContext
public:
 
   /// Accesses the context reference.
   PyMlirContextRef &getContext() { return contextRef; }
 
 private:
   PyMlirContextRef contextRef;
 };
 
+/// Wrapper around an MlirLocation.
+class PyLocation : public BaseContextObject {
+public:
+  PyLocation(PyMlirContextRef contextRef, MlirLocation loc)
+      : BaseContextObject(std::move(contextRef)), loc(loc) {}
+
+  operator MlirLocation() const { return loc; }
+  MlirLocation get() const { return loc; }
+
+  /// Enter and exit the context manager.
+  pybind11::object contextEnter();
+  void contextExit(const pybind11::object &excType,
+                   const pybind11::object &excVal,
+                   const pybind11::object &excTb);
+
+  /// Gets a capsule wrapping the void* within the MlirLocation.
+  pybind11::object getCapsule();
+
+  /// Creates a PyLocation from the MlirLocation wrapped by a capsule.
+  /// Note that PyLocation instances are uniqued, so the returned object
+  /// may be a pre-existing object. Ownership of the underlying MlirLocation
+  /// is taken by calling this function.
+  static PyLocation createFromCapsule(pybind11::object capsule);
+
+private:
+  MlirLocation loc;
+};
+
 /// Python class mirroring the C MlirDiagnostic struct. Note that these structs
 /// are only valid for the duration of a diagnostic callback and attempting
 /// to access them outside of that will raise an exception. This applies to
 /// nested diagnostics (in the notes) as well.
 class PyDiagnostic {
 public:
   PyDiagnostic(MlirDiagnostic diagnostic) : diagnostic(diagnostic) {}
   void invalidate();
   bool isValid() { return valid; }
   MlirDiagnosticSeverity getSeverity();
   PyLocation getLocation();
   pybind11::str getMessage();
   pybind11::tuple getNotes();
 
+  /// Materialized diagnostic information. This is safe to access outside the
+  /// diagnostic callback.
+  struct DiagnosticInfo {
+    MlirDiagnosticSeverity severity;
+    PyLocation location;
+    std::string message;
+    std::vector<DiagnosticInfo> notes;
+  };
+  DiagnosticInfo getInfo();
+
 private:
   MlirDiagnostic diagnostic;
 
   void checkValid();
   /// If notes have been materialized from the diagnostic, then this will
   /// be populated with the corresponding objects (all castable to
   /// PyDiagnostic).
   std::optional<pybind11::tuple> materializedNotes;

 private:
   MlirContext context;
   pybind11::object callback;
   std::optional<MlirDiagnosticHandlerID> registeredID;
   bool hadError = false;
   friend class PyMlirContext;
 };
 
+/// RAII object that captures any error diagnostics emitted to the provided
+/// context.
+struct PyMlirContext::ErrorCapture {
+  ErrorCapture(PyMlirContextRef ctx)
+      : ctx(ctx), handlerID(mlirContextAttachDiagnosticHandler(
+                      ctx->get(), handler, /*userData=*/this,
+                      /*deleteUserData=*/nullptr)) {}
+  ~ErrorCapture() {
+    mlirContextDetachDiagnosticHandler(ctx->get(), handlerID);
+    assert(errors.empty() && "unhandled captured errors");
+  }
+
+  std::vector<PyDiagnostic::DiagnosticInfo> take() {
+    return std::move(errors);
+  };
+
+private:
+  PyMlirContextRef ctx;
+  MlirDiagnosticHandlerID handlerID;
+  std::vector<PyDiagnostic::DiagnosticInfo> errors;
+
+  static MlirLogicalResult handler(MlirDiagnostic diag, void *userData);
+};
+
 /// Wrapper around an MlirDialect. This is exported as `DialectDescriptor` in
 /// order to differentiate it from the `Dialect` base class which is extended by
 /// plugins which extend dialect functionality through extension python code.
 /// This should be seen as the "low-level" object and `Dialect` as the
 /// high-level, user facing object.
 class PyDialectDescriptor : public BaseContextObject {
 public:
   PyDialectDescriptor(PyMlirContextRef contextRef, MlirDialect dialect)
public:
 
   pybind11::object getCapsule();
   static PyDialectRegistry createFromCapsule(pybind11::object capsule);
 
 private:
   MlirDialectRegistry registry;
 };
 
-/// Wrapper around an MlirLocation.
-class PyLocation : public BaseContextObject {
-public:
-  PyLocation(PyMlirContextRef contextRef, MlirLocation loc)
-      : BaseContextObject(std::move(contextRef)), loc(loc) {}
-
-  operator MlirLocation() const { return loc; }
-  MlirLocation get() const { return loc; }
-
-  /// Enter and exit the context manager.
-  pybind11::object contextEnter();
-  void contextExit(const pybind11::object &excType,
-                   const pybind11::object &excVal,
-                   const pybind11::object &excTb);
-
-  /// Gets a capsule wrapping the void* within the MlirLocation.
-  pybind11::object getCapsule();
-
-  /// Creates a PyLocation from the MlirLocation wrapped by a capsule.
-  /// Note that PyLocation instances are uniqued, so the returned object
-  /// may be a pre-existing object. Ownership of the underlying MlirLocation
-  /// is taken by calling this function.
-  static PyLocation createFromCapsule(pybind11::object capsule);
-
-private:
-  MlirLocation loc;
-};
-
 /// Used in function arguments when None should resolve to the current context
 /// manager set instance.
 class DefaultingPyLocation
     : public Defaulting<DefaultingPyLocation, PyLocation> {
 public:
   using Defaulting::Defaulting;
   static constexpr const char kTypeDescription[] = "mlir.ir.Location";
   static PyLocation &resolve();
public:
 
   // Implement the bound 'writeBytecode' method.
   void writeBytecode(const pybind11::object &fileObject);
 
   /// Moves the operation before or after the other operation.
   void moveAfter(PyOperationBase &other);
   void moveBefore(PyOperationBase &other);
 
+  /// Verify the operation. Throws `MLIRError` if verification fails, and
+  /// returns `true` otherwise.
+  bool verify();
+
   /// Each must provide access to the raw Operation.
   virtual PyOperation &getOperation() = 0;
 };
 
 /// Wrapper around PyOperation.
 /// Operations exist in either an attached (dependent) or detached (top-level)
 /// state. In the detached state (as on creation), an operation is owned by
 /// the creator and its lifetime extends either until its reference count
public:
   /// Casts the bindings class into the C API structure.
   operator MlirSymbolTable() { return symbolTable; }
 
 private:
   PyOperationRef operation;
   MlirSymbolTable symbolTable;
 };
 
+/// Custom exception that allows access to error diagnostic information. This is
+/// converted to the `ir.MLIRError` python exception when thrown.
+struct MLIRError {
+  MLIRError(llvm::Twine message,
+            std::vector<PyDiagnostic::DiagnosticInfo> &&errorDiagnostics = {})
+      : message(message.str()), errorDiagnostics(std::move(errorDiagnostics)) {}
+  std::string message;
+  std::vector<PyDiagnostic::DiagnosticInfo> errorDiagnostics;
+};
+
 void populateIRAffine(pybind11::module &m);
 void populateIRAttributes(pybind11::module &m);
 void populateIRCore(pybind11::module &m);
 void populateIRInterfaces(pybind11::module &m);
 void populateIRTypes(pybind11::module &m);
 
 } // namespace python
 } // namespace mlir

mlir/lib/Bindings/Python/IRTypes.cpp

public:
   static constexpr const char *pyClassName = "VectorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            DefaultingPyLocation loc) {
+          PyMlirContext::ErrorCapture errors(loc->getContext());
           MlirType t = mlirVectorTypeGetChecked(loc, shape.size(), shape.data(),
                                                 elementType);
-          // TODO: Rework error reporting once diagnostic engine is exposed
-          // in C API.
-          if (mlirTypeIsNull(t)) {
-            throw SetPyError(
-                PyExc_ValueError,
-                Twine("invalid '") +
-                    py::repr(py::cast(elementType)).cast<std::string>() +
-                    "' and expected floating point or integer type.");
-          }
+          if (mlirTypeIsNull(t))
+            throw MLIRError("Invalid type", errors.take());
           return PyVectorType(elementType.getContext(), t);
         },
         py::arg("shape"), py::arg("elementType"), py::arg("loc") = py::none(),
         "Create a vector type");
   }
 };
 
 /// Ranked Tensor Type subclass - RankedTensorType.
 class PyRankedTensorType
     : public PyConcreteType<PyRankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
   static constexpr const char *pyClassName = "RankedTensorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            std::optional<PyAttribute> &encodingAttr, DefaultingPyLocation loc) {
+          PyMlirContext::ErrorCapture errors(loc->getContext());
           MlirType t = mlirRankedTensorTypeGetChecked(
               loc, shape.size(), shape.data(), elementType,
               encodingAttr ? encodingAttr->get() : mlirAttributeGetNull());
-          // TODO: Rework error reporting once diagnostic engine is exposed
-          // in C API.
-          if (mlirTypeIsNull(t)) {
-            throw SetPyError(
-                PyExc_ValueError,
-                Twine("invalid '") +
-                    py::repr(py::cast(elementType)).cast<std::string>() +
-                    "' and expected floating point, integer, vector or "
-                    "complex "
-                    "type.");
-          }
+          if (mlirTypeIsNull(t))
+            throw MLIRError("Invalid type", errors.take());
           return PyRankedTensorType(elementType.getContext(), t);
         },
         py::arg("shape"), py::arg("element_type"),
         py::arg("encoding") = py::none(), py::arg("loc") = py::none(),
         "Create a ranked tensor type");
     c.def_property_readonly(
         "encoding", [](PyRankedTensorType &self) -> std::optional<PyAttribute> {
           MlirAttribute encoding = mlirRankedTensorTypeGetEncoding(self.get());
public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedTensor;
   static constexpr const char *pyClassName = "UnrankedTensorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &elementType, DefaultingPyLocation loc) {
+          PyMlirContext::ErrorCapture errors(loc->getContext());
           MlirType t = mlirUnrankedTensorTypeGetChecked(loc, elementType);
-          // TODO: Rework error reporting once diagnostic engine is exposed
-          // in C API.
-          if (mlirTypeIsNull(t)) {
-            throw SetPyError(
-                PyExc_ValueError,
-                Twine("invalid '") +
-                    py::repr(py::cast(elementType)).cast<std::string>() +
-                    "' and expected floating point, integer, vector or "
-                    "complex "
-                    "type.");
-          }
+          if (mlirTypeIsNull(t))
+            throw MLIRError("Invalid type", errors.take());
           return PyUnrankedTensorType(elementType.getContext(), t);
         },
         py::arg("element_type"), py::arg("loc") = py::none(),
         "Create a unranked tensor type");
   }
 };
 
 /// Ranked MemRef Type subclass - MemRefType.
 class PyMemRefType : public PyConcreteType<PyMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAMemRef;
   static constexpr const char *pyClassName = "MemRefType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](std::vector<int64_t> shape, PyType &elementType,
             PyAttribute *layout, PyAttribute *memorySpace,
             DefaultingPyLocation loc) {
+           PyMlirContext::ErrorCapture errors(loc->getContext());
            MlirAttribute layoutAttr = layout ? *layout : mlirAttributeGetNull();
            MlirAttribute memSpaceAttr =
                memorySpace ? *memorySpace : mlirAttributeGetNull();
            MlirType t =
                mlirMemRefTypeGetChecked(loc, elementType, shape.size(),
                                         shape.data(), layoutAttr, memSpaceAttr);
-           // TODO: Rework error reporting once diagnostic engine is exposed
-           // in C API.
-           if (mlirTypeIsNull(t)) {
-             throw SetPyError(
-                 PyExc_ValueError,
-                 Twine("invalid '") +
-                     py::repr(py::cast(elementType)).cast<std::string>() +
-                     "' and expected floating point, integer, vector or "
-                     "complex "
-                     "type.");
-           }
+           if (mlirTypeIsNull(t))
+             throw MLIRError("Invalid type", errors.take());
            return PyMemRefType(elementType.getContext(), t);
          },
          py::arg("shape"), py::arg("element_type"),
          py::arg("layout") = py::none(), py::arg("memory_space") = py::none(),
          py::arg("loc") = py::none(), "Create a memref type")
         .def_property_readonly(
             "layout",
             [](PyMemRefType &self) -> PyAttribute {
public:
   static constexpr const char *pyClassName = "UnrankedMemRefType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](PyType &elementType, PyAttribute *memorySpace,
             DefaultingPyLocation loc) {
+           PyMlirContext::ErrorCapture errors(loc->getContext());
            MlirAttribute memSpaceAttr = {};
            if (memorySpace)
              memSpaceAttr = *memorySpace;
 
            MlirType t =
                mlirUnrankedMemRefTypeGetChecked(loc, elementType, memSpaceAttr);
-           // TODO: Rework error reporting once diagnostic engine is exposed
-           // in C API.
-           if (mlirTypeIsNull(t)) {
-             throw SetPyError(
-                 PyExc_ValueError,
-                 Twine("invalid '") +
-                     py::repr(py::cast(elementType)).cast<std::string>() +
-                     "' and expected floating point, integer, vector or "
-                     "complex "
-                     "type.");
-           }
+           if (mlirTypeIsNull(t))
+             throw MLIRError("Invalid type", errors.take());
            return PyUnrankedMemRefType(elementType.getContext(), t);
          },
          py::arg("element_type"), py::arg("memory_space"),
          py::arg("loc") = py::none(), "Create a unranked memref type")
         .def_property_readonly(
             "memory_space",
             [](PyUnrankedMemRefType &self) -> PyAttribute {
               MlirAttribute a = mlirMemRefTypeGetMemorySpace(self);

mlir/lib/Bindings/Python/Pass.cpp

py::class_<PyPassManager>(m, "PassManager", py::module_local())
               throw SetPyError(PyExc_ValueError, std::string(errorMsg.join()));
           },
           py::arg("pipeline"),
           "Add textual pipeline elements to the pass manager. Throws a "
           "ValueError if the pipeline can't be parsed.")
       .def(
           "run",
           [](PyPassManager &passManager, PyOperationBase &op) {
+            PyMlirContext::ErrorCapture errors(op.getOperation().getContext());
             MlirLogicalResult status = mlirPassManagerRunOnOp(
                 passManager.get(), op.getOperation().get());
             if (mlirLogicalResultIsFailure(status))
-              throw SetPyError(PyExc_RuntimeError,
-                               "Failure while executing pass pipeline.");
+              throw MLIRError("Failure while executing pass pipeline",
+                              errors.take());
           },
           py::arg("operation"),
-          "Run the pass manager on the provided operation, throw a "
-          "RuntimeError on failure.")
+          "Run the pass manager on the provided operation, raising an "
+          "MLIRError on failure.")
       .def(
           "__str__",
           [](PyPassManager &self) {
             MlirPassManager passManager = self.get();
             PyPrintAccumulator printAccum;
             mlirPrintPassPipeline(
                 mlirPassManagerGetAsOpPassManager(passManager),
                 printAccum.getCallback(), printAccum.getUserData());
             return printAccum.join();
           },
           "Print the textual representation for this PassManager, suitable to "
           "be passed to `parse` for round-tripping.");
 }

mlir/python/mlir/_mlir_libs/__init__.py

def __init__(self, *args, **kwargs):
       super().__init__(*args, **kwargs)
       self.append_dialect_registry(registry)
       for hook in post_init_hooks:
         hook(self)
       # TODO: There is some debate about whether we should eagerly load
       # all dialects. It is being done here in order to preserve existing
       # behavior. See: https://github.com/llvm/llvm-project/issues/56037
       self.load_all_available_dialects()
-
   ir.Context = Context
 
+  class MLIRError(Exception):
+    """
+    An exception with diagnostic information. Has the following fields:
+      message: str
+      error_diagnostics: List[ir.DiagnosticInfo]
+    """
+    def __init__(self, message, error_diagnostics):
+      self.message = message
+      self.error_diagnostics = error_diagnostics
+      super().__init__(message, error_diagnostics)
+
+    def __str__(self):
+      s = self.message
+      if self.error_diagnostics:
+        s += ':'
+      for diag in self.error_diagnostics:
+        s += "\nerror: "  + str(diag.location)[4:-1] + ": " + diag.message.replace('\n', '\n  ')
+        for note in diag.notes:
+          s += "\n note: "  + str(note.location)[4:-1] + ": " + note.message.replace('\n', '\n  ')
+      return s
+  ir.MLIRError = MLIRError
+
 
 _site_initialize()

mlir/test/python/ir/attributes.py

def testParsePrint():
   gc.collect()
   # CHECK: "hello"
   print(str(t))
   # CHECK: Attribute("hello")
   print(repr(t))
 
 
 # CHECK-LABEL: TEST: testParseError
-# TODO: Hook the diagnostic manager to capture a more meaningful error
-# message.
 @run
 def testParseError():
   with Context():
     try:
       t = Attribute.parse("BAD_ATTR_DOES_NOT_EXIST")
-    except ValueError as e:
-      # CHECK: Unable to parse attribute: 'BAD_ATTR_DOES_NOT_EXIST'
-      print("testParseError:", e)
+    except MLIRError as e:
+      # CHECK: testParseError: <
+      # CHECK:   Unable to parse attribute:
+      # CHECK:   error: "BAD_ATTR_DOES_NOT_EXIST":1:1: expected attribute value
+      # CHECK: >
+      print(f"testParseError: <{e}>")
     else:
       print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testAttrEq
 @run
 def testAttrEq():
   with Context():
print("default_get:", FloatAttr.get(
         F32Type.get(), 42.0))
     # CHECK: f32_get: 4.200000e+01 : f32
     print("f32_get:", FloatAttr.get_f32(42.0))
     # CHECK: f64_get: 4.200000e+01 : f64
     print("f64_get:", FloatAttr.get_f64(42.0))
     try:
       fattr_invalid = FloatAttr.get(
           IntegerType.get_signless(32), 42)
-    except ValueError as e:
-      # CHECK: invalid 'Type(i32)' and expected floating point type.
+    except MLIRError as e:
+      # CHECK: Invalid attribute:
+      # CHECK: error: unknown: expected floating point type
       print(e)
     else:
       print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testIntegerAttr
 @run
 def testIntegerAttr():

mlir/test/python/ir/builtin_types.py

def testParsePrint():
   gc.collect()
   # CHECK: i32
   print(str(t))
   # CHECK: Type(i32)
   print(repr(t))
 
 
 # CHECK-LABEL: TEST: testParseError
-# TODO: Hook the diagnostic manager to capture a more meaningful error
-# message.
 @run
 def testParseError():
   ctx = Context()
   try:
     t = Type.parse("BAD_TYPE_DOES_NOT_EXIST", ctx)
-  except ValueError as e:
-    # CHECK: Unable to parse type: 'BAD_TYPE_DOES_NOT_EXIST'
-    print("testParseError:", e)
+  except MLIRError as e:
+    # CHECK: testParseError: <
+    # CHECK:   Unable to parse type:
+    # CHECK:   error: "BAD_TYPE_DOES_NOT_EXIST":1:1: expected non-function type
+    # CHECK: >
+    print(f"testParseError: <{e}>")
   else:
     print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testTypeEq
 @run
 def testTypeEq():
   ctx = Context()
with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     # CHECK: vector type: vector<2x3xf32>
     print("vector type:", VectorType.get(shape, f32))
 
     none = NoneType.get()
     try:
       vector_invalid = VectorType.get(shape, none)
-    except ValueError as e:
-      # CHECK: invalid 'Type(none)' and expected floating point or integer type.
+    except MLIRError as e:
+      # CHECK: Invalid type:
+      # CHECK: error: unknown: vector elements must be int/index/float type but got 'none'
       print(e)
     else:
       print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testRankedTensorType
 @run
 def testRankedTensorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     loc = Location.unknown()
     # CHECK: ranked tensor type: tensor<2x3xf32>
     print("ranked tensor type:",
           RankedTensorType.get(shape, f32))
 
     none = NoneType.get()
     try:
       tensor_invalid = RankedTensorType.get(shape, none)
-    except ValueError as e:
-      # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
-      # CHECK: or complex type.
+    except MLIRError as e:
+      # CHECK: Invalid type:
+      # CHECK: error: unknown: invalid tensor element type: 'none'
       print(e)
     else:
       print("Exception not produced")
 
     # Encoding should be None.
     assert RankedTensorType.get(shape, f32).encoding is None
 
     tensor = RankedTensorType.get(shape, f32)
except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
 
     none = NoneType.get()
     try:
       tensor_invalid = UnrankedTensorType.get(none)
-    except ValueError as e:
-      # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
-      # CHECK: or complex type.
+    except MLIRError as e:
+      # CHECK: Invalid type:
+      # CHECK: error: unknown: invalid tensor element type: 'none'
       print(e)
     else:
       print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testMemRefType
 @run
 def testMemRefType():
with Context(), Location.unknown():
     # CHECK: memref affine map: (d0, d1) -> (d1, d0)
     print("memref affine map:", memref_layout.affine_map)
     # CHECK: memory space: <<NULL ATTRIBUTE>>
     print("memory space:", memref_layout.memory_space)
 
     none = NoneType.get()
     try:
       memref_invalid = MemRefType.get(shape, none)
-    except ValueError as e:
-      # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
-      # CHECK: or complex type.
+    except MLIRError as e:
+      # CHECK: Invalid type:
+      # CHECK: error: unknown: invalid memref element type
       print(e)
     else:
       print("Exception not produced")
 
     assert memref.shape == shape
 
 
 # CHECK-LABEL: TEST: testUnrankedMemRefType
except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
 
     none = NoneType.get()
     try:
       memref_invalid = UnrankedMemRefType.get(none, Attribute.parse("2"))
-    except ValueError as e:
-      # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
-      # CHECK: or complex type.
+    except MLIRError as e:
+      # CHECK: Invalid type:
+      # CHECK: error: unknown: invalid memref element type
       print(e)
     else:
       print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testTupleType
 @run
 def testTupleType():

mlir/test/python/ir/diagnostic_handler.py

def testDiagnosticEmptyNotes():
   loc.emit_error("foobar")
   assert not handler.had_error
 
 
 # CHECK-LABEL: TEST: testDiagnosticNonEmptyNotes
 @run
 def testDiagnosticNonEmptyNotes():
   ctx = Context()
+  ctx.emit_error_diagnostics = True
   def callback(d):
     # CHECK: DIAGNOSTIC:
     # CHECK:   message='arith.addi' op requires one result
     # CHECK:   notes=['see current operation: "arith.addi"() : () -> ()']
     print(f"DIAGNOSTIC:")
     print(f"  message={d.message}")
     print(f"  notes={list(map(str, d.notes))}")
     return True
   handler = ctx.attach_diagnostic_handler(callback)
   loc = Location.unknown(ctx)
+  try:
     Operation.create('arith.addi', loc=loc).verify()
+  except MLIRError:
+    pass
   assert not handler.had_error
 
 # CHECK-LABEL: TEST: testDiagnosticCallbackException
 @run
 def testDiagnosticCallbackException():
   ctx = Context()
   def callback(d):
     raise ValueError("Error in handler")

mlir/test/python/ir/exception.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import gc
+from mlir.ir import *
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+  return f
+
+
+# CHECK-LABEL: TEST: test_exception
+@run
+def test_exception():
+  ctx =  Context()
+  ctx.allow_unregistered_dialects = True
+  try:
+    Operation.parse("""
+      func.func @foo() {
+          "test.use"(%0) : (i64) -> ()  loc("use")
+          %0 = "test.def"() : () -> i64 loc("def")
+          return
+      }
+    """, context=ctx)
+  except MLIRError as e:
+    # CHECK: Exception: <
+    # CHECK:   Unable to parse operation assembly:
+    # CHECK:   error: "use": operand #0 does not dominate this use
+    # CHECK:    note: "use": see current operation: "test.use"(%0) : (i64) -> ()
+    # CHECK:    note: "def": operand defined here (op in the same block)
+    # CHECK: >
+    print(f"Exception: <{e}>")
+
+    # CHECK: message: Unable to parse operation assembly
+    print(f"message: {e.message}")
+
+    # CHECK: error_diagnostics[0]:           loc("use") operand #0 does not dominate this use
+    # CHECK: error_diagnostics[0].notes[0]:  loc("use") see current operation: "test.use"(%0) : (i64) -> ()
+    # CHECK: error_diagnostics[0].notes[1]:  loc("def") operand defined here (op in the same block)
+    print("error_diagnostics[0]:          ", e.error_diagnostics[0].location, e.error_diagnostics[0].message)
+    print("error_diagnostics[0].notes[0]: ", e.error_diagnostics[0].notes[0].location, e.error_diagnostics[0].notes[0].message)
+    print("error_diagnostics[0].notes[1]: ", e.error_diagnostics[0].notes[1].location, e.error_diagnostics[0].notes[1].message)
+
+
+# CHECK-LABEL: test_emit_error_diagnostics
+@run
+def test_emit_error_diagnostics():
+  ctx = Context()
+  loc = Location.unknown(ctx)
+  handler_diags = []
+  def handler(d):
+    handler_diags.append(str(d))
+    return True
+  ctx.attach_diagnostic_handler(handler)
+
+  try:
+    Attribute.parse("not an attr", ctx)
+  except MLIRError as e:
+    # CHECK: emit_error_diagnostics=False:
+    # CHECK: e.error_diagnostics: ['expected attribute value']
+    # CHECK: handler_diags: []
+    print(f"emit_error_diagnostics=False:")
+    print(f"e.error_diagnostics: {[str(diag) for diag in e.error_diagnostics]}")
+    print(f"handler_diags: {handler_diags}")
+
+  ctx.emit_error_diagnostics = True
+  try:
+    Attribute.parse("not an attr", ctx)
+  except MLIRError as e:
+    # CHECK: emit_error_diagnostics=True:
+    # CHECK: e.error_diagnostics: []
+    # CHECK: handler_diags: ['expected attribute value']
+    print(f"emit_error_diagnostics=True:")
+    print(f"e.error_diagnostics: {[str(diag) for diag in e.error_diagnostics]}")
+    print(f"handler_diags: {handler_diags}")

mlir/test/python/ir/module.py

def testParseSuccess():
   ctx = None  # Ensure that module captures the context.
   gc.collect()
   module.dump()  # Just outputs to stderr. Verifies that it functions.
   print(str(module))
 
 
 # Verify parse error.
 # CHECK-LABEL: TEST: testParseError
-# CHECK: testParseError: Unable to parse module assembly (see diagnostics)
+# CHECK: testParseError: <
+# CHECK:   Unable to parse module assembly:
+# CHECK:   error: "-":1:1: expected operation name in quotes
+# CHECK: >
 @run
 def testParseError():
   ctx = Context()
   try:
     module = Module.parse(r"""}SYNTAX ERROR{""", ctx)
-  except ValueError as e:
-    print("testParseError:", e)
+  except MLIRError as e:
+    print(f"testParseError: <{e}>")
   else:
     print("Exception not produced")
 
 
 # Verify successful parse.
 # CHECK-LABEL: TEST: testCreateEmpty
 # CHECK: module {
 @run

mlir/test/python/ir/operation.py

 def testInvalidOperationStrSoftFails():
   ctx = Context()
   with Location.unknown(ctx):
     invalid_op = create_invalid_operation()
     # Verify that we fallback to the generic printer for safety.
     # CHECK: "builtin.module"() ({
     # CHECK: }) : () -> ()
     print(invalid_op)
-    # CHECK: .verify = False
-    print(f".verify = {invalid_op.operation.verify()}")
+    try:
+      invalid_op.verify()
+    except MLIRError as e:
+      # CHECK: Exception: <
+      # CHECK:   Verification failed:
+      # CHECK:   error: unknown: 'builtin.module' op requires one region
+      # CHECK:    note: unknown: see current operation:
+      # CHECK:     "builtin.module"() ({
+      # CHECK:     ^bb0:
+      # CHECK:     }, {
+      # CHECK:     }) : () -> ()
+      # CHECK: >
+      print(f"Exception: <{e}>")
 
 
 # CHECK-LABEL: TEST: testInvalidModuleStrSoftFails
 @run
 def testInvalidModuleStrSoftFails():
   ctx = Context()
   with Location.unknown(ctx):
     module = Module.create()
with Context() as ctx:
     assert isinstance(m, ModuleOp)
     assert type(o) is OpView
 
     # Parsing specific operation.
     m = ModuleOp.parse('module {}')
     assert isinstance(m, ModuleOp)
     try:
       ModuleOp.parse('"test.foo"() : () -> ()')
-    except ValueError as e:
+    except MLIRError as e:
       # CHECK: error: Expected a 'builtin.module' op, got: 'test.foo'
       print(f"error: {e}")
     else:
       assert False, "expected error"
 
     o = Operation.parse('"test.foo"() : () -> ()', source_name="my-source-string")
     # CHECK: op_with_source_name: "test.foo"() : () -> () loc("my-source-string":1:1)
     print(f"op_with_source_name: {o.get_asm(enable_debug_info=True, use_local_scope=True)}")

mlir/test/python/pass_manager.py

except ValueError as e:
       # CHECK: ValueError exception: Can't add pass 'NormalizeMemRefs' restricted to 'builtin.module' on a PassManager intended to run on 'func.func', did you intend to nest?
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 run(testInvalidNesting)
 
 
 # Verify that a pass manager can execute on IR
-# CHECK-LABEL: TEST: testRun
+# CHECK-LABEL: TEST: testRunPipeline
 def testRunPipeline():
   with Context():
     pm = PassManager.parse("any(print-op-stats{json=false})")
     func = FuncOp.parse(r"""func.func @successfulParse() { return }""")
     pm.run(func)
 # CHECK: Operations encountered:
 # CHECK: func.func      , 1
 # CHECK: func.return        , 1
 run(testRunPipeline)
+
+# CHECK-LABEL: TEST: testRunPipelineError
+@run
+def testRunPipelineError():
+  with Context() as ctx:
+    ctx.allow_unregistered_dialects = True
+    op = Operation.parse('"test.op"() : () -> ()')
+    pm = PassManager.parse("any(cse)")
+    try:
+      pm.run(op)
+    except MLIRError as e:
+      # CHECK: Exception: <
+      # CHECK:   Failure while executing pass pipeline:
+      # CHECK:   error: "-":1:1: 'test.op' op trying to schedule a pass on an unregistered operation
+      # CHECK:    note: "-":1:1: see current operation: "test.op"() : () -> ()
+      # CHECK: >
+      print(f"Exception: <{e}>")