[Mlir-commits] [mlir] 04d76d3 - [mlir][python] Add nameloc getter

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2021-10-12T12:45:57-07:00
New Revision: 04d76d36948cf298914155f787d5f7d2e2264624

URL: https://github.com/llvm/llvm-project/commit/04d76d36948cf298914155f787d5f7d2e2264624
DIFF: https://github.com/llvm/llvm-project/commit/04d76d36948cf298914155f787d5f7d2e2264624.diff

LOG: [mlir][python] Add nameloc getter

Expose the nameloc getter to Python API.

Differential Revision: https://reviews.llvm.org/D111663

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/test/python/ir/location.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8ed3bd5ed38a8..3226b7fe180f5 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -43,6 +43,9 @@ See also: https://mlir.llvm.org/docs/LangRef/#type-system
 static const char kContextGetFileLocationDocstring[] =
     R"(Gets a Location representing a file, line and column)";
 
+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.
 
@@ -1970,6 +1973,19 @@ void mlir::python::populateIRCore(py::module &m) {
           },
           py::arg("filename"), py::arg("line"), py::arg("col"),
           py::arg("context") = py::none(), kContextGetFileLocationDocstring)
+      .def_static(
+          "name",
+          [](std::string name, llvm::Optional<PyLocation> childLoc,
+             DefaultingPyMlirContext context) {
+            return PyLocation(
+                context->getRef(),
+                mlirLocationNameGet(
+                    context->get(), toMlirStringRef(name),
+                    childLoc ? childLoc->get()
+                             : 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(); },

diff  --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index 42a96a1ba604e..b17e5ec6be6a8 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -39,6 +39,25 @@ def testFileLineCol():
 run(testFileLineCol)
 
 
+# CHECK-LABEL: TEST: testName
+def testName():
+  with Context() as ctx:
+    loc = Location.name("nombre")
+    locWithChildLoc = Location.name("naam", loc)
+  ctx = None
+  gc.collect()
+  # CHECK: file str: loc("nombre")
+  print("file str:", str(loc))
+  # CHECK: file repr: loc("nombre")
+  print("file repr:", repr(loc))
+  # CHECK: file str: loc("naam"("nombre"))
+  print("file str:", str(locWithChildLoc))
+  # CHECK: file repr: loc("naam"("nombre"))
+  print("file repr:", repr(locWithChildLoc))
+
+run(testName)
+
+
 # CHECK-LABEL: TEST: testLocationCapsule
 def testLocationCapsule():
   with Context() as ctx: