[Mlir-commits] [mlir] [mlir][Python] Downcast location returned from diagnostic (PR #201337)

[Martin Erhart]

https://github.com/maerhart created https://github.com/llvm/llvm-project/pull/201337

Essentially a follow up to https://github.com/llvm/llvm-project/pull/192630

>From 44c8ccc7a5303975c48cc97066008bf8ca0c624e Mon Sep 17 00:00:00 2001
From: Martin Erhart <martin.erhart at sifive.com>
Date: Wed, 3 Jun 2026 13:04:47 +0100
Subject: [PATCH] [mlir][Python] Downcast location returned from diagnostic

---
 mlir/include/mlir/Bindings/Python/IRCore.h |  2 +-
 mlir/lib/Bindings/Python/IRCore.cpp        |  8 ++++----
 mlir/test/python/ir/location.py            | 23 ++++++++++++++++++++++
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 0b758d4061d67..1b8c232b009c0 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -375,7 +375,7 @@ class MLIR_PYTHON_API_EXPORTED PyDiagnostic {
   void invalidate();
   bool isValid() { return valid; }
   PyDiagnosticSeverity getSeverity();
-  PyLocation getLocation();
+  nanobind::typed<nanobind::object, PyLocation> getLocation();
   nanobind::str getMessage();
   nanobind::typed<nanobind::tuple, PyDiagnostic> getNotes();
 
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index f92d4c14ceb16..92e9ecf3f2c20 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -754,11 +754,11 @@ PyDiagnosticSeverity PyDiagnostic::getSeverity() {
       mlirDiagnosticGetSeverity(diagnostic));
 }
 
-PyLocation PyDiagnostic::getLocation() {
+nb::typed<nb::object, PyLocation> PyDiagnostic::getLocation() {
   checkValid();
   MlirLocation loc = mlirDiagnosticGetLocation(diagnostic);
   MlirContext context = mlirLocationGetContext(loc);
-  return PyLocation(PyMlirContext::forContext(context), loc);
+  return PyLocation(PyMlirContext::forContext(context), loc).maybeDownCast();
 }
 
 nb::str PyDiagnostic::getMessage() {
@@ -789,8 +789,8 @@ PyDiagnostic::DiagnosticInfo PyDiagnostic::getInfo() {
   std::vector<DiagnosticInfo> notes;
   for (nb::handle n : getNotes())
     notes.emplace_back(nb::cast<PyDiagnostic>(n).getInfo());
-  return {getSeverity(), getLocation(), nb::cast<std::string>(getMessage()),
-          std::move(notes)};
+  return {getSeverity(), nb::cast<PyLocation>(getLocation()),
+          nb::cast<std::string>(getMessage()), std::move(notes)};
 }
 
 //------------------------------------------------------------------------------
diff --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index 33a4ffce48b9c..579aa9fa7c7c3 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -292,3 +292,26 @@ def testLocationCapsule():
 
 
 run(testLocationCapsule)
+
+
+# CHECK-LABEL: TEST: testLocationFromDiagnostic
+def testLocationFromDiagnostic():
+    with Context() as ctx:
+        def callback(d):
+            assert isinstance(d.location, FileLineColLoc)
+            assert isinstance(d.location, Location)
+
+            # CHECK: filename: diagnostic_test.txt
+            print("filename:", d.location.filename)
+            # CHECK: line: 42
+            print("line:", d.location.start_line)
+            # CHECK: col: 7
+            print("col:", d.location.start_col)
+            return True
+
+        ctx.attach_diagnostic_handler(callback)
+        loc = FileLineColLoc.get("diagnostic_test.txt", 42, 7)
+        loc.emit_error("test error message")
+
+
+run(testLocationFromDiagnostic)