| 15 | } |
| 16 | |
| 17 | nb::object tree_map( |
| 18 | const std::vector<nb::object>& trees, |
| 19 | std::function<nb::object(const std::vector<nb::object>&)> transform) { |
| 20 | std::function<nb::object(const std::vector<nb::object>&)> recurse; |
| 21 | |
| 22 | recurse = [&](const std::vector<nb::object>& subtrees) { |
| 23 | if (nb::isinstance<nb::list>(subtrees[0])) { |
| 24 | nb::list l; |
| 25 | std::vector<nb::object> items(subtrees.size()); |
| 26 | validate_subtrees<nb::list, nb::tuple, nb::dict>(subtrees); |
| 27 | for (int i = 0; i < nb::cast<nb::list>(subtrees[0]).size(); ++i) { |
| 28 | for (int j = 0; j < subtrees.size(); ++j) { |
| 29 | if (nb::isinstance<nb::list>(subtrees[j])) { |
| 30 | items[j] = nb::cast<nb::list>(subtrees[j])[i]; |
| 31 | } else { |
| 32 | items[j] = subtrees[j]; |
| 33 | } |
| 34 | } |
| 35 | l.append(recurse(items)); |
| 36 | } |
| 37 | return nb::cast<nb::object>(l); |
| 38 | } else if (nb::isinstance<nb::tuple>(subtrees[0])) { |
| 39 | // Check the rest of the subtrees |
| 40 | std::vector<nb::object> items(subtrees.size()); |
| 41 | int len = nb::cast<nb::tuple>(subtrees[0]).size(); |
| 42 | nb::list l; |
| 43 | validate_subtrees<nb::tuple, nb::list, nb::dict>(subtrees); |
| 44 | auto type = subtrees[0].type(); |
| 45 | for (int i = 0; i < len; ++i) { |
| 46 | for (int j = 0; j < subtrees.size(); ++j) { |
| 47 | if (nb::isinstance<nb::tuple>(subtrees[j])) { |
| 48 | items[j] = nb::cast<nb::tuple>(subtrees[j])[i]; |
| 49 | } else { |
| 50 | items[j] = subtrees[j]; |
| 51 | } |
| 52 | } |
| 53 | l.append(recurse(items)); |
| 54 | } |
| 55 | if (PyTuple_CheckExact(subtrees[0].ptr())) { |
| 56 | return nb::cast<nb::object>(nb::tuple(l)); |
| 57 | } |
| 58 | return nb::hasattr(type, "_fields") ? type(*l) : type(l); |
| 59 | } else if (nb::isinstance<nb::dict>(subtrees[0])) { |
| 60 | std::vector<nb::object> items(subtrees.size()); |
| 61 | validate_subtrees<nb::dict, nb::list, nb::tuple>(subtrees); |
| 62 | nb::dict d; |
| 63 | for (auto item : nb::cast<nb::dict>(subtrees[0])) { |
| 64 | for (int j = 0; j < subtrees.size(); ++j) { |
| 65 | if (nb::isinstance<nb::dict>(subtrees[j])) { |
| 66 | auto subdict = nb::cast<nb::dict>(subtrees[j]); |
| 67 | if (!subdict.contains(item.first)) { |
| 68 | throw std::invalid_argument( |
| 69 | "[tree_map] Tree is not a valid prefix tree of the first tree."); |
| 70 | } |
| 71 | items[j] = subdict[item.first]; |
| 72 | } else { |
| 73 | items[j] = subtrees[j]; |
| 74 | } |
no test coverage detected