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

[Maksim Levental]

================
@@ -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);
+            }
+          },
+          "callback"_a, "op_class"_a = nb::none(),
+          "walk_order"_a = PyWalkOrder::PostOrder,
+          // clang-format off
+           nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], op_class: type[OpView] | None = None, walk_order: WalkOrder = WalkOrder.POST_ORDER) -> None"),
----------------
makslevental wrote:

can you double check the stubs generaetd from this - for whatever reason stubgen on `WalkOrder` malfunctions (you can actually take a look here https://github.com/llvm/llvm-project/pull/186106/changes#diff-acfa9a5ea4439eae3f3e5a57ccbb8e72944842f6ae7816624c1d75e2908f3e62R3933 to see how @superbobry is currently changing/updating that stub)

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