[Mlir-commits] [mlir] [MLIR][Python] Add type filter to walk() binding and add get_ops_of_type() utility (PR #186131)

[Jakub Kuderski]

================
@@ -3858,16 +3858,41 @@ void populateIRCore(nb::module_ &m) {
 
             Note:
               After erasing, any Python references to the operation become invalid.)")
-      .def("walk", &PyOperationBase::walk, "callback"_a,
-           "walk_order"_a = PyWalkOrder::PostOrder,
-           // clang-format off
-          nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder) -> None"),
-           // clang-format on
-           R"(
+      .def(
+          "walk",
+          [](PyOperationBase &self,
+             std::function<PyWalkResult(MlirOperation)> callback,
+             std::optional<nb::object> opClass, PyWalkOrder walkOrder) {
+            if (opClass) {
+              self.walk(
+                  [&](MlirOperation mlirOp) -> PyWalkResult {
+                    nb::object opview =
+                        PyOperation::forOperation(
+                            self.getOperation().getContext(), mlirOp)
+                            ->createOpView();
+                    if (nb::isinstance(opview, *opClass))
+                      return callback(mlirOp);
+                    return PyWalkResult::Advance;
+                  },
+                  walkOrder);
+            } else {
+              self.walk(callback, walkOrder);
+            }
----------------
kuhar wrote:

Flip the if condition above and handle this first, then return early, and alst do what we have in the `if` body. See https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code

https://github.com/llvm/llvm-project/pull/186131