| 171 | } |
| 172 | |
| 173 | void tree_visit_update( |
| 174 | nb::object tree, |
| 175 | std::function<nb::object(nb::handle)> visitor) { |
| 176 | std::function<nb::object(nb::handle)> recurse; |
| 177 | recurse = [&](nb::handle subtree) { |
| 178 | if (nb::isinstance<nb::list>(subtree)) { |
| 179 | auto l = nb::cast<nb::list>(subtree); |
| 180 | for (int i = 0; i < l.size(); ++i) { |
| 181 | l[i] = recurse(l[i]); |
| 182 | } |
| 183 | return nb::cast<nb::object>(l); |
| 184 | } else if (nb::isinstance<nb::tuple>(subtree)) { |
| 185 | auto type = subtree.type(); |
| 186 | nb::list l(subtree); |
| 187 | for (int i = 0; i < l.size(); ++i) { |
| 188 | l[i] = recurse(l[i]); |
| 189 | } |
| 190 | if (PyTuple_CheckExact(subtree.ptr())) { |
| 191 | return nb::cast<nb::object>(nb::tuple(l)); |
| 192 | } |
| 193 | return nb::hasattr(type, "_fields") ? type(*l) : type(l); |
| 194 | } else if (nb::isinstance<nb::dict>(subtree)) { |
| 195 | auto d = nb::cast<nb::dict>(subtree); |
| 196 | for (auto item : d) { |
| 197 | d[item.first] = recurse(item.second); |
| 198 | } |
| 199 | return nb::cast<nb::object>(d); |
| 200 | } else if (nb::isinstance<mx::array>(subtree)) { |
| 201 | return visitor(subtree); |
| 202 | } else { |
| 203 | return nb::cast<nb::object>(subtree); |
| 204 | } |
| 205 | }; |
| 206 | recurse(tree); |
| 207 | } |
| 208 | |
| 209 | // Fill a pytree (recursive dict or list of dict or list) |
| 210 | // in place with the given arrays |
no test coverage detected