Adding a __iter__ binding to Dims

Signed-off-by: Yinghai Lu <yinghai@fb.com>
NVIDIA · Dec 6, 2021 · 34a1675 · 34a1675
1 parent 6f38570
commit 34a1675
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/src/infer/pyFoundationalTypes.cpp b/python/src/infer/pyFoundationalTypes.cpp
@@ -239,6 +239,9 @@ void bindFoundationalTypes(py::module& m)
         .def("__eq__", lambdas::dimsEqual<Dims, py::tuple>)
         // These functions allow us to use Dims like an iterable.
         .def("__len__", lambdas::dims_len)
+        .def("__iter__", [](const Dims &s) {
+              return py::make_iterator(&s.d[0], &s.d[s.nbDims]); },
+              py::keep_alive<0, 1>())
         .def("__getitem__", lambdas::dims_getter)
         .def("__getitem__", lambdas::dims_getter_slice)
         .def("__setitem__", lambdas::dims_setter)