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

[rkayaith]

Summary

Add a read-only loc property to Operation and OpView

[ftynse]

Thanks! A couple of naming nits.

mlir/include/mlir-c/IR.h

361

We tend to use full words in C API (and in Python below), so mlirOperationGetLocation.

mlir/lib/Bindings/Python/IRCore.cpp

2147

"location" please.

2149

Please expand auto. We use auto when the type is obvious from context or excessively long to spell out.

[rkayaith]

address comments

[ftynse]

Please let me know if you need help landing this.

[rkayaith]

Thanks for the review, could you land this?

mlir/include/mlir-c/IR.h

 /// Checks whether two operation handles point to the same operation. This does
 /// not perform deep comparison.
 MLIR_CAPI_EXPORTED bool mlirOperationEqual(MlirOperation op,
                                            MlirOperation other);
 
 /// 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 mlirOperationGetLoc(MlirOperation op);
+
 /// Gets the type id of the operation.
 /// Returns null if the operation does not have a registered operation
 /// description.
 MLIR_CAPI_EXPORTED MlirTypeID mlirOperationGetTypeID(MlirOperation op);
 
 /// Gets the name of the operation as an identifier.
 MLIR_CAPI_EXPORTED MlirIdentifier mlirOperationGetName(MlirOperation op);
 

mlir/lib/Bindings/Python/IRCore.cpp

py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
                       " results (it is only valid for operations with a "
                       "single result)");
             }
             return PyOpResult(operation.getRef(),
                               mlirOperationGetResult(operation, 0));
           },
           "Shortcut to get an op result if it has only one (throws an error "
           "otherwise).")
+      .def_property_readonly(
+          "loc",
+          [](PyOperationBase &self) {
+            auto &operation = self.getOperation();
+            return PyLocation(operation.getContext(),
+                              mlirOperationGetLoc(operation.get()));
+          },
+          "Returns the source location the operation was defined or derived "
+          "from.")
       .def("__iter__",
            [](PyOperationBase &self) {
              return PyRegionIterator(self.getOperation().getRef());
            })
       .def(
           "__str__",
           [](PyOperationBase &self) {
             return self.getAsm(/*binary=*/false,

mlir/lib/CAPI/IR/IR.cpp

 bool mlirOperationEqual(MlirOperation op, MlirOperation other) {
   return unwrap(op) == unwrap(other);
 }
 
 MlirContext mlirOperationGetContext(MlirOperation op) {
   return wrap(unwrap(op)->getContext());
 }
 
+MlirLocation mlirOperationGetLoc(MlirOperation op) {
+  return wrap(unwrap(op)->getLoc());
+}
+
 MlirTypeID mlirOperationGetTypeID(MlirOperation op) {
   if (const auto *abstractOp = unwrap(op)->getAbstractOperation()) {
     return wrap(abstractOp->typeID);
   }
   return {nullptr};
 }
 
 MlirIdentifier mlirOperationGetName(MlirOperation op) {

mlir/test/python/ir/operation.py

with InsertionPoint(m.body):
 
       op.operation.erase()
 
       # CHECK-NOT: "custom.op1"
       print(m)
 
       # 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.loc == loc
+    assert op.operation.loc == loc