Add __iter__ to RDom bindings (#8861)

* Add __iter__ to RDom Python bindings

Previously, iterating over an RDom with list() or a for loop would fail
with a KeyError. While __len__ and __getitem__ were defined, the Python
iterator protocol requires __iter__ to return an iterator object.

This adds a lightweight RDomIterator helper and exposes it via __iter__,
allowing natural iteration:
  r = hl.RDom([(0, 10), (0, 20)])
  for v in r: ...
  list(r)  # returns [RVar, RVar]

Adds test coverage in the existing rdom.py correctness test.

* Fix Python test label propagation in AddPythonTest.cmake

The LABEL argument was not being propagated to set_tests_properties,
causing -L filters to not work correctly. Now both 'python' and the
specific label (e.g., 'python_correctness') are applied to each test.

Co-Authored-By: GitHub Copilot <github-copilot@users.noreply.github.com>

Author: alexreinking

python_bindings/cmake/AddPythonTest.cmake

@@ -13,14 +13,16 @@ function(add_python_test)
     cmake_path(GET ARG_FILE STEM test_name)
     set(test_name "${ARG_LABEL}_${test_name}")
 
+    set(test_labels python ${ARG_LABEL})
+
     add_test(
         NAME "${test_name}"
         COMMAND ${Halide_PYTHON_LAUNCHER} "$<TARGET_FILE:Python::Interpreter>" "$<SHELL_PATH:${CMAKE_CURRENT_SOURCE_DIR}/${ARG_FILE}>" ${ARG_TEST_ARGS}
     )
     set_tests_properties(
         "${test_name}"
         PROPERTIES
-        LABELS "python"
+        LABELS "${test_labels}"
         ENVIRONMENT "${ARG_ENVIRONMENT}"
         ENVIRONMENT_MODIFICATION "${ARG_PYTHONPATH}"
         SKIP_REGULAR_EXPRESSION "\\[SKIP\\]"

python_bindings/src/halide/halide_/PyRDom.cpp

@@ -30,6 +30,33 @@ void define_rvar(py::module &m) {
 void define_rdom(py::module &m) {
     define_rvar(m);
 
+    // A small iterator wrapper to expose RDom iteration to Python.
+    // It holds a copy of the RDom and an index, and implements
+    // the iterator protocol (__iter__ and __next__).
+    struct RDomIterator {
+        RDom rd;
+        int idx = 0;
+        RDomIterator() = default;
+        RDomIterator(const RDom &r) : rd(r) {
+        }
+        RVar next() {
+            if (idx >= rd.dimensions()) {
+                throw py::stop_iteration();
+            }
+            return rd[idx++];
+        }
+        RDomIterator &iter() {
+            return *this;
+        }
+    };
+
+    // Expose the iterator type to Python so we can return it from __iter__.
+    py::class_<RDomIterator>(m, "_RDomIterator")
+        .def(py::init<>())
+        .def(py::init<const RDom &>())
+        .def("__iter__", &RDomIterator::iter)
+        .def("__next__", &RDomIterator::next);
+
     auto rdom_class =
         py::class_<RDom>(m, "RDom")
             .def(py::init<>())
@@ -41,12 +68,13 @@ void define_rdom(py::module &m) {
             .def("same_as", &RDom::same_as)
             .def("dimensions", &RDom::dimensions)
             .def("__len__", &RDom::dimensions)
+            .def("__iter__", [](const RDom &r) { return RDomIterator(r); }, py::keep_alive<0, 1>())
             .def("where", &RDom::where, py::arg("predicate"))
             .def("__getitem__", [](RDom &r, const int i) -> RVar {
                 if (i < 0 || i >= r.dimensions()) {
                     throw pybind11::key_error();
                 }
-                return r[i];
+                return r[i];  //
             })
             .def_readonly("x", &RDom::x)
             .def_readonly("y", &RDom::y)
@@ -55,7 +83,7 @@ void define_rdom(py::module &m) {
             .def("__repr__", [](const RDom &r) -> std::string {
                 std::ostringstream o;
                 o << "<halide.RDom " << r << ">";
-                return o.str();
+                return o.str();  //
             });
 
     add_binary_operators(rdom_class);

python_bindings/test/correctness/rdom.py

@@ -46,6 +46,16 @@ def test_rdom():
     return 0
 
 
+def test_rdom_iter():
+    # Verify that RDom is iterable and yields RVars matching index access
+    r = hl.RDom([(0, 10), (0, 20)], "r")
+    it = list(r)
+    assert len(it) == len(r)
+    assert all(hasattr(v, "min") and hasattr(v, "extent") for v in it)
+    for i in range(len(r)):
+        assert str(it[i]) == str(r[i])
+
+
 def test_implicit_pure_definition():
     a = np.random.ranf((2, 3)).astype(np.float32)
     expected = a.sum(axis=1)
@@ -67,3 +77,4 @@ def test_implicit_pure_definition():
 if __name__ == "__main__":
     test_rdom()
     test_implicit_pure_definition()
+    test_rdom_iter()