[mlir][python] Add bindings for diagnostic handler.

[stellaraccident]

Summary

I considered multiple approaches for this but settled on this one because I could make the lifetime management work in a reasonably easy way (others had issues with not being able to cast to a Python reference from a C++ constructor). We could stand to have more formatting helpers, but best to get the core mechanism in first.

[stellaraccident]

Update type metadata.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModule.h

336

When are the contextEnter and contextExit method called?

[stellaraccident]

mlir/lib/Bindings/Python/IRModule.h

336

They are bound to __enter__ and __exit__ to enable use as a context manager. We do similar naming elsewhere in this module.

[mehdi_amini]

I'd rather wait for @ftynse to have a look as well, I think I understand everything in this patch, but this in the realm of "I wouldn't be able to write this myself" ;)

mlir/lib/Bindings/Python/IRModule.h

336

Right, I wasn't clear, I should have asked: "when is this class used as a context manager?"

But I figured in the meantime by reading the test, so it's fine :)

[stellaraccident]

In D116568#3218814, @mehdi_amini wrote:

I'd rather wait for @ftynse to have a look as well, I think I understand everything in this patch, but this in the realm of "I wouldn't be able to write this myself" ;)

No problem at all getting another set of eyes... I've been needing/dreading this for the past year. A few more days makes no difference :)

re: "I wouldn't be able to write this myself" - I'm sure you could if sufficiently motivated.

For Alex, the primary alternative design I considered/tried was having the DiagnosticHandler be a free-standing class that auto attached on construction like:

with Context(), DiagnosticHandler(callback):
  # do stuff

Trying to make this work ran into all kinds of undocumented and surprising pybind behavior, since you need to get access to the py::object from within the scope of the py::init<> function (so that you can borrow against the python reference count). I couldn't come up with a way to make ownership transfer properly in that situation. I don't expect this is going to be one of those everyday things for people to be doing and opted for a simpler implementation. My hope is that we layer some kind of exception handler thing on top of this so that I can do something like:

with CaptureDiagnostics() as diag_cap:
  # run passes and do stuff
  # If we exit the with exceptionally, then raise a chained exception that includes any diagnostics.

# If we exit successfully:
print(diag_cap.diagnostics)

(or something like that -- need to fiddle with it, and this can be written in pure Python atop this basic facility)

[ftynse]

LGTM! I agree that it's better to write nicer APIs directly in Python, debugging Pybind glue is hard.

mlir/lib/Bindings/Python/IRCore.cpp

515

Nit: this could use a comment that the Python will free the newed object during its GC sweep because of ownership transfer below.

mlir/test/python/ir/diagnostic_handler.py

75

An optional kwargs notes? As in loc.emit_error("foo", notes=[loc.error_note("bar"), other_loc.error_note("qux")]) ?

[stellaraccident]

mlir/test/python/ir/diagnostic_handler.py

75

Yes, something like that. Will need more C-APIs (and more Python), which was further than I wanted to go in this patch.

[stellaraccident]

Comment.

mlir/lib/Bindings/Python/IRCore.cpp

 }
 
 void PyMlirContext::contextExit(const pybind11::object &excType,
                                 const pybind11::object &excVal,
                                 const pybind11::object &excTb) {
   PyThreadContextEntry::popContext(*this);
 }
 
+py::object PyMlirContext::attachDiagnosticHandler(py::object callback) {
+  // Note that ownership is transferred to the delete callback below by way of
+  // an explicit inc_ref (borrow).
+  PyDiagnosticHandler *pyHandler =
+      new PyDiagnosticHandler(get(), std::move(callback));
+  py::object pyHandlerObject =
+      py::cast(pyHandler, py::return_value_policy::take_ownership);
+  pyHandlerObject.inc_ref();
+
+  // In these C callbacks, the userData is a PyDiagnosticHandler* that is
+  // guaranteed to be known to pybind.
+  auto handlerCallback =
+      +[](MlirDiagnostic diagnostic, void *userData) -> MlirLogicalResult {
+    PyDiagnostic *pyDiagnostic = new PyDiagnostic(diagnostic);
+    py::object pyDiagnosticObject =
+        py::cast(pyDiagnostic, py::return_value_policy::take_ownership);
+
+    auto *pyHandler = static_cast<PyDiagnosticHandler *>(userData);
+    bool result = false;
+    {
+      // Since this can be called from arbitrary C++ contexts, always get the
+      // gil.
+      py::gil_scoped_acquire gil;
+      try {
+        result = py::cast<bool>(pyHandler->callback(pyDiagnostic));
+      } catch (std::exception &e) {
+        fprintf(stderr, "MLIR Python Diagnostic handler raised exception: %s\n",
+                e.what());
+        pyHandler->hadError = true;
+      }
+    }
+
+    pyDiagnostic->invalidate();
+    return result ? mlirLogicalResultSuccess() : mlirLogicalResultFailure();
+  };
+  auto deleteCallback = +[](void *userData) {
+    auto *pyHandler = static_cast<PyDiagnosticHandler *>(userData);
+    assert(pyHandler->registeredID && "handler is not registered");
+    pyHandler->registeredID.reset();
+
+    // Decrement reference, balancing the inc_ref() above.
+    py::object pyHandlerObject =
+        py::cast(pyHandler, py::return_value_policy::reference);
+    pyHandlerObject.dec_ref();
+  };
+
+  pyHandler->registeredID = mlirContextAttachDiagnosticHandler(
+      get(), handlerCallback, static_cast<void *>(pyHandler), deleteCallback);
+  return pyHandlerObject;
+}
+
 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():'");
if (stack.empty())
     throw SetPyError(PyExc_RuntimeError, "Unbalanced Location enter/exit");
   auto &tos = stack.back();
   if (tos.frameKind != FrameKind::Location && tos.getLocation() != &location)
     throw SetPyError(PyExc_RuntimeError, "Unbalanced Location enter/exit");
   stack.pop_back();
 }
 
 //------------------------------------------------------------------------------
+// PyDiagnostic*
+//------------------------------------------------------------------------------
+
+void PyDiagnostic::invalidate() {
+  valid = false;
+  if (materializedNotes) {
+    for (auto &noteObject : *materializedNotes) {
+      PyDiagnostic *note = py::cast<PyDiagnostic *>(noteObject);
+      note->invalidate();
+    }
+  }
+}
+
+PyDiagnosticHandler::PyDiagnosticHandler(MlirContext context,
+                                         py::object callback)
+    : context(context), callback(std::move(callback)) {}
+
+PyDiagnosticHandler::~PyDiagnosticHandler() {}
+
+void PyDiagnosticHandler::detach() {
+  if (!registeredID)
+    return;
+  MlirDiagnosticHandlerID localID = *registeredID;
+  mlirContextDetachDiagnosticHandler(context, localID);
+  assert(!registeredID && "should have unregistered");
+  // Not strictly necessary but keeps stale pointers from being around to cause
+  // issues.
+  context = {nullptr};
+}
+
+void PyDiagnostic::checkValid() {
+  if (!valid) {
+    throw std::invalid_argument(
+        "Diagnostic is invalid (used outside of callback)");
+  }
+}
+
+MlirDiagnosticSeverity PyDiagnostic::getSeverity() {
+  checkValid();
+  return mlirDiagnosticGetSeverity(diagnostic);
+}
+
+PyLocation PyDiagnostic::getLocation() {
+  checkValid();
+  MlirLocation loc = mlirDiagnosticGetLocation(diagnostic);
+  MlirContext context = mlirLocationGetContext(loc);
+  return PyLocation(PyMlirContext::forContext(context), loc);
+}
+
+py::str PyDiagnostic::getMessage() {
+  checkValid();
+  py::object fileObject = py::module::import("io").attr("StringIO")();
+  PyFileAccumulator accum(fileObject, /*binary=*/false);
+  mlirDiagnosticPrint(diagnostic, accum.getCallback(), accum.getUserData());
+  return fileObject.attr("getvalue")();
+}
+
+py::tuple PyDiagnostic::getNotes() {
+  checkValid();
+  if (materializedNotes)
+    return *materializedNotes;
+  intptr_t numNotes = mlirDiagnosticGetNumNotes(diagnostic);
+  materializedNotes = py::tuple(numNotes);
+  for (intptr_t i = 0; i < numNotes; ++i) {
+    MlirDiagnostic noteDiag = mlirDiagnosticGetNote(diagnostic, i);
+    py::object pyNoteDiag = py::cast(PyDiagnostic(noteDiag));
+    PyTuple_SET_ITEM(materializedNotes->ptr(), i, pyNoteDiag.ptr());
+  }
+  return *materializedNotes;
+}
+
+//------------------------------------------------------------------------------
 // PyDialect, PyDialectDescriptor, PyDialects
 //------------------------------------------------------------------------------
 
 MlirDialect PyDialects::getDialectForKey(const std::string &key,
                                          bool attrError) {
   MlirDialect dialect = mlirContextGetOrLoadDialect(getContext()->get(),
                                                     {key.data(), key.size()});
   if (mlirDialectIsNull(dialect)) {

 } // namespace
 
 //------------------------------------------------------------------------------
 // Populates the core exports of the 'ir' submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRCore(py::module &m) {
   //----------------------------------------------------------------------------
+  // Enums.
+  //----------------------------------------------------------------------------
+  py::enum_<MlirDiagnosticSeverity>(m, "DiagnosticSeverity", py::module_local())
+      .value("ERROR", MlirDiagnosticError)
+      .value("WARNING", MlirDiagnosticWarning)
+      .value("NOTE", MlirDiagnosticNote)
+      .value("REMARK", MlirDiagnosticRemark);
+
+  //----------------------------------------------------------------------------
+  // Mapping of Diagnostics.
+  //----------------------------------------------------------------------------
+  py::class_<PyDiagnostic>(m, "Diagnostic", py::module_local())
+      .def_property_readonly("severity", &PyDiagnostic::getSeverity)
+      .def_property_readonly("location", &PyDiagnostic::getLocation)
+      .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_<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);
+
+  //----------------------------------------------------------------------------
   // Mapping of MlirContext.
   //----------------------------------------------------------------------------
   py::class_<PyMlirContext>(m, "Context", py::module_local())
       .def(py::init<>(&PyMlirContext::createNewContextForInit))
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
            [](PyMlirContext &self) {
              PyMlirContextRef ref = PyMlirContext::forContext(self.get());
py::class_<PyMlirContext>(m, "Context", py::module_local())
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
           },
           [](PyMlirContext &self, bool value) {
             mlirContextSetAllowUnregisteredDialects(self.get(), value);
           })
+      .def("attach_diagnostic_handler", &PyMlirContext::attachDiagnosticHandler,
+           py::arg("callback"),
+           "Attaches a diagnostic handler that will receive callbacks")
       .def(
           "enable_multithreading",
           [](PyMlirContext &self, bool enable) {
             mlirContextEnableMultithreading(self.get(), enable);
           },
           py::arg("enable"))
       .def(
           "is_registered_operation",
py::class_<PyLocation>(m, "Location", py::module_local())
                 context->getRef(),
                 mlirLocationFileLineColGet(
                     context->get(), toMlirStringRef(filename), line, col));
           },
           py::arg("filename"), py::arg("line"), py::arg("col"),
           py::arg("context") = py::none(), kContextGetFileLocationDocstring)
       .def_static(
           "fused",
-          [](const std::vector<PyLocation> &pyLocations, llvm::Optional<PyAttribute> metadata,
+          [](const std::vector<PyLocation> &pyLocations,
+             llvm::Optional<PyAttribute> metadata,
              DefaultingPyMlirContext context) {
             if (pyLocations.empty())
               throw py::value_error("No locations provided");
             llvm::SmallVector<MlirLocation, 4> locations;
             locations.reserve(pyLocations.size());
             for (auto &pyLocation : pyLocations)
               locations.push_back(pyLocation.get());
             MlirLocation location = mlirLocationFusedGet(
py::class_<PyLocation>(m, "Location", py::module_local())
                              : mlirLocationUnknownGet(context->get())));
           },
           py::arg("name"), py::arg("childLoc") = py::none(),
           py::arg("context") = py::none(), kContextGetNameLocationDocString)
       .def_property_readonly(
           "context",
           [](PyLocation &self) { return self.getContext().getObject(); },
           "Context that owns the Location")
+      .def(
+          "emit_error",
+          [](PyLocation &self, std::string message) {
+            mlirEmitError(self, message.c_str());
+          },
+          py::arg("message"), "Emits an error at this location")
       .def("__repr__", [](PyLocation &self) {
         PyPrintAccumulator printAccum;
         mlirLocationPrint(self, printAccum.getCallback(),
                           printAccum.getUserData());
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------

mlir/lib/Bindings/Python/IRModule.h

 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <vector>
 
 #include "PybindUtils.h"
 
 #include "mlir-c/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
+#include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/IntegerSet.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 
 namespace mlir {
 namespace python {
 
 class PyBlock;
+class PyDiagnostic;
+class PyDiagnosticHandler;
 class PyInsertionPoint;
 class PyLocation;
 class DefaultingPyLocation;
 class PyMlirContext;
 class DefaultingPyMlirContext;
 class PyModule;
 class PyOperation;
 class PyType;
public:
   size_t getLiveModuleCount();
 
   /// Enter and exit the context manager.
   pybind11::object contextEnter();
   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);
+
 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.
public:
 
   /// Accesses the context reference.
   PyMlirContextRef &getContext() { return contextRef; }
 
 private:
   PyMlirContextRef contextRef;
 };
 
+/// 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();
+
+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).
+  llvm::Optional<pybind11::tuple> materializedNotes;
+  bool valid = true;
+};
+
+/// Represents a diagnostic handler attached to the context. The handler's
+/// callback will be invoked with PyDiagnostic instances until the detach()
+/// method is called or the context is destroyed. A diagnostic handler can be
+/// the subject of a `with` block, which will detach it when the block exits.
+///
+/// Since diagnostic handlers can call back into Python code which can do
+/// unsafe things (i.e. recursively emitting diagnostics, raising exceptions,
+/// etc), this is generally not deemed to be a great user-level API. Users
+/// should generally use some form of DiagnosticCollector. If the handler raises
+/// any exceptions, they will just be emitted to stderr and dropped.
+///
+/// The unique usage of this class means that its lifetime management is
+/// different from most other parts of the API. Instances are always created
+/// in an attached state and can transition to a detached state by either:
+///   a) The context being destroyed and unregistering all handlers.
+///   b) An explicit call to detach().
+/// The object may remain live from a Python perspective for an arbitrary time
+/// after detachment, but there is nothing the user can do with it (since there
+/// is no way to attach an existing handler object).
+class PyDiagnosticHandler {
+public:
+  PyDiagnosticHandler(MlirContext context, pybind11::object callback);
+  ~PyDiagnosticHandler();
+
+  bool isAttached() { return registeredID.hasValue(); }
+  bool getHadError() { return hadError; }
+
+  /// Detaches the handler. Does nothing if not attached.
+  void detach();
+
+  pybind11::object contextEnter() { return pybind11::cast(this); }
+  void contextExit(pybind11::object excType, pybind11::object excVal,
+                   pybind11::object excTb) {
+    detach();
+  }
+
+private:
+  MlirContext context;
+  pybind11::object callback;
+  llvm::Optional<MlirDiagnosticHandlerID> registeredID;
+  bool hadError = false;
+  friend class PyMlirContext;
+};
+
 /// 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)

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

 # Originally imported via:
 #   stubgen {...} -m mlir._mlir_libs._mlir.ir
 # Local modifications:
 #   * Rewrite references to 'mlir.ir.' to local types
 #   * Add __all__ with the following incantation:
 #       egrep '^class ' ir.pyi | awk -F ' |:|\\(' '{print "    \"" $2 "\","}'
 #   * Local edits to signatures and types that MyPy did not auto detect (or
 #     detected incorrectly).
 
-from typing import Any, Callable, ClassVar, Dict, List, Optional, Sequence
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Sequence, Tuple
 
 from typing import overload
 
 __all__ = [
     "AffineAddExpr",
     "AffineBinaryExpr",
     "AffineCeilDivExpr",
     "AffineConstantExpr",
__all__ = [
     "ComplexType",
     "Context",
     "DenseElementsAttr",
     "DenseFPElementsAttr",
     "DenseIntElementsAttr",
     "Dialect",
     "DialectDescriptor",
     "Dialects",
+    "Diagnostic",
+    "DiagnosticHandler",
+    "DiagnosticSeverity",
     "DictAttr",
     "F16Type",
     "F32Type",
     "F64Type",
     "FlatSymbolRefAttr",
     "FloatAttr",
     "FunctionType",
     "IndexType",
class Context:
     allow_unregistered_dialects: bool
     def __init__(self) -> None: ...
     def _CAPICreate(self) -> object: ...
     def _get_context_again(self) -> "Context": ...
     @staticmethod
     def _get_live_count() -> int: ...
     def _get_live_module_count(self) -> int: ...
     def _get_live_operation_count(self) -> int: ...
+    def attach_diagnostic_handler(self, callback: Callable[["Diagnostic"], bool]) -> "DiagnosticHandler": ...
     def enable_multithreading(self, enable: bool) -> None: ...
-    def get_dialect_descriptor(name: dialect_name: str) -> "DialectDescriptor": ...
+    def get_dialect_descriptor(dialect_name: str) -> "DialectDescriptor": ...
     def is_registered_operation(self, operation_name: str) -> bool: ...
     def __enter__(self) -> "Context": ...
     def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
     @property
     def _CAPIPtr(self) -> object: ...
     @property
     def d(self) -> "Dialects": ...
     @property
class DialectDescriptor:
     @property
     def namespace(self) -> str: ...
 
 class Dialects:
     def __init__(self, *args, **kwargs) -> None: ...
     def __getattr__(self, arg0: str) -> "Dialect": ...
     def __getitem__(self, arg0: str) -> "Dialect": ...
 
+class Diagnostic:
+    @property
+    def severity(self) -> "DiagnosticSeverity": ...
+    @property
+    def location(self) -> "Location": ...
+    @property
+    def message(self) -> str: ...
+    @property
+    def notes(self) -> Tuple["Diagnostic"]: ...
+
+class DiagnosticHandler:
+    def detach(self) -> None: ...
+    @property
+    def attached(self) -> bool: ...
+    @property
+    def had_error(self) -> bool: ...
+    def __enter__(self) -> "DiagnosticHandler": ...
+    def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
+
+class DiagnosticSeverity:
+    ERROR: "DiagnosticSeverity"
+    WARNING: "DiagnosticSeverity"
+    NOTE: "DiagnosticSeverity"
+    REMARK: "DiagnosticSeverity"
+
 # TODO: Auto-generated. Audit and fix.
 class DictAttr(Attribute):
     def __init__(self, cast_from_attr: Attribute) -> None: ...
     def get(self, *args, **kwargs) -> Any: ...
     def isinstance(self, *args, **kwargs) -> Any: ...
     def __contains__(self, arg0: str) -> bool: ...
     @overload
     def __getitem__(self, arg0: str) -> Attribute: ...

mlir/test/python/ir/diagnostic_handler.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
+
+
+@run
+def testLifecycleContextDestroy():
+  ctx = Context()
+  def callback(foo): ...
+  handler = ctx.attach_diagnostic_handler(callback)
+  assert handler.attached
+  # If context is destroyed before the handler, it should auto-detach.
+  ctx = None
+  gc.collect()
+  assert not handler.attached
+
+  # And finally collecting the handler should be fine.
+  handler = None
+  gc.collect()
+
+
+@run
+def testLifecycleExplicitDetach():
+  ctx = Context()
+  def callback(foo): ...
+  handler = ctx.attach_diagnostic_handler(callback)
+  assert handler.attached
+  handler.detach()
+  assert not handler.attached
+
+
+@run
+def testLifecycleWith():
+  ctx = Context()
+  def callback(foo): ...
+  with ctx.attach_diagnostic_handler(callback) as handler:
+    assert handler.attached
+  assert not handler.attached
+
+
+@run
+def testLifecycleWithAndExplicitDetach():
+  ctx = Context()
+  def callback(foo): ...
+  with ctx.attach_diagnostic_handler(callback) as handler:
+    assert handler.attached
+    handler.detach()
+  assert not handler.attached
+
+
+# CHECK-LABEL: TEST: testDiagnosticCallback
+@run
+def testDiagnosticCallback():
+  ctx = Context()
+  def callback(d):
+    # CHECK: DIAGNOSTIC: message='foobar', severity=DiagnosticSeverity.ERROR, loc=loc(unknown)
+    print(f"DIAGNOSTIC: message='{d.message}', severity={d.severity}, loc={d.location}")
+    return True
+  handler = ctx.attach_diagnostic_handler(callback)
+  loc = Location.unknown(ctx)
+  loc.emit_error("foobar")
+  assert not handler.had_error
+
+
+# CHECK-LABEL: TEST: testDiagnosticEmptyNotes
+# TODO: Come up with a way to inject a diagnostic with notes from this API.
+@run
+def testDiagnosticEmptyNotes():
+  ctx = Context()
+  def callback(d):
+    # CHECK: DIAGNOSTIC: notes=()
+    print(f"DIAGNOSTIC: notes={d.notes}")
+    return True
+  handler = ctx.attach_diagnostic_handler(callback)
+  loc = Location.unknown(ctx)
+  loc.emit_error("foobar")
+  assert not handler.had_error
+
+
+# CHECK-LABEL: TEST: testDiagnosticCallbackException
+@run
+def testDiagnosticCallbackException():
+  ctx = Context()
+  def callback(d):
+    raise ValueError("Error in handler")
+  handler = ctx.attach_diagnostic_handler(callback)
+  loc = Location.unknown(ctx)
+  loc.emit_error("foobar")
+  assert handler.had_error
+
+
+# CHECK-LABEL: TEST: testEscapingDiagnostic
+@run
+def testEscapingDiagnostic():
+  ctx = Context()
+  diags = []
+  def callback(d):
+    diags.append(d)
+    return True
+  handler = ctx.attach_diagnostic_handler(callback)
+  loc = Location.unknown(ctx)
+  loc.emit_error("foobar")
+  assert not handler.had_error
+
+  # CHECK: DIAGNOSTIC: <Invalid Diagnostic>
+  print(f"DIAGNOSTIC: {str(diags[0])}")
+  try:
+    diags[0].severity
+    raise RuntimeError("expected exception")
+  except ValueError:
+    pass
+  try:
+    diags[0].location
+    raise RuntimeError("expected exception")
+  except ValueError:
+    pass
+  try:
+    diags[0].message
+    raise RuntimeError("expected exception")
+  except ValueError:
+    pass
+  try:
+    diags[0].notes
+    raise RuntimeError("expected exception")
+  except ValueError:
+    pass
+
+
+
+# CHECK-LABEL: TEST: testDiagnosticReturnTrueHandles
+@run
+def testDiagnosticReturnTrueHandles():
+  ctx = Context()
+  def callback1(d):
+    print(f"CALLBACK1: {d}")
+    return True
+  def callback2(d):
+    print(f"CALLBACK2: {d}")
+    return True
+  ctx.attach_diagnostic_handler(callback1)
+  ctx.attach_diagnostic_handler(callback2)
+  loc = Location.unknown(ctx)
+  # CHECK-NOT: CALLBACK1
+  # CHECK: CALLBACK2: foobar
+  # CHECK-NOT: CALLBACK1
+  loc.emit_error("foobar")
+
+
+# CHECK-LABEL: TEST: testDiagnosticReturnFalseDoesNotHandle
+@run
+def testDiagnosticReturnFalseDoesNotHandle():
+  ctx = Context()
+  def callback1(d):
+    print(f"CALLBACK1: {d}")
+    return True
+  def callback2(d):
+    print(f"CALLBACK2: {d}")
+    return False
+  ctx.attach_diagnostic_handler(callback1)
+  ctx.attach_diagnostic_handler(callback2)
+  loc = Location.unknown(ctx)
+  # CHECK: CALLBACK2: foobar
+  # CHECK: CALLBACK1: foobar
+  loc.emit_error("foobar")