[mlir][python] Add 'loc' property to ops

Add a read-only `loc` property to Operation and OpView

Reviewed By: ftynse

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

Author: rkayaith

mlir/include/mlir-c/IR.h

@@ -357,6 +357,9 @@ MLIR_CAPI_EXPORTED bool mlirOperationEqual(MlirOperation op,
 /// Gets the context this operation is associated with
 MLIR_CAPI_EXPORTED MlirContext mlirOperationGetContext(MlirOperation op);
 
+/// Gets the location of the operation.
+MLIR_CAPI_EXPORTED MlirLocation mlirOperationGetLocation(MlirOperation op);
+
 /// Gets the type id of the operation.
 /// Returns null if the operation does not have a registered operation
 /// description.

mlir/lib/Bindings/Python/IRCore.cpp

@@ -2143,6 +2143,15 @@ void mlir::python::populateIRCore(py::module &m) {
           },
           "Shortcut to get an op result if it has only one (throws an error "
           "otherwise).")
+      .def_property_readonly(
+          "location",
+          [](PyOperationBase &self) {
+            PyOperation &operation = self.getOperation();
+            return PyLocation(operation.getContext(),
+                              mlirOperationGetLocation(operation.get()));
+          },
+          "Returns the source location the operation was defined or derived "
+          "from.")
       .def("__iter__",
            [](PyOperationBase &self) {
              return PyRegionIterator(self.getOperation().getRef());

mlir/lib/CAPI/IR/IR.cpp

@@ -346,6 +346,10 @@ MlirContext mlirOperationGetContext(MlirOperation op) {
   return wrap(unwrap(op)->getContext());
 }
 
+MlirLocation mlirOperationGetLocation(MlirOperation op) {
+  return wrap(unwrap(op)->getLoc());
+}
+
 MlirTypeID mlirOperationGetTypeID(MlirOperation op) {
   if (const auto *abstractOp = unwrap(op)->getAbstractOperation()) {
     return wrap(abstractOp->typeID);

mlir/test/python/ir/operation.py

@@ -728,3 +728,15 @@ def testOperationErase():
 
       # Ensure we can create another operation
       Operation.create("custom.op2")
+
+
+# CHECK-LABEL: TEST: testOperationLoc
+@run
+def testOperationLoc():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with ctx:
+    loc = Location.name("loc")
+    op = Operation.create("custom.op", loc=loc)
+    assert op.location == loc
+    assert op.operation.location == loc