Tree tools

Tree tools#

Some tools for working with python trees, useful for working with nested data structures.

icomo.walk_tree(tree)[source]#

Walk through a tree and yield the indices and values of the leaves.

Parameters:

tree (PyTree) – The tree to walk through. Can only be a nested tree of lists, tuples and/or dictionaries.

Yields:

indices, value – The indices as a list and values of the leaves.

Return type:

Generator[tuple[list, Any]]

icomo.nested_indexing(tree, indices, add=None, at=None)[source]#

Return the element of a nested structure of lists or tuples.

Parameters:
  • tree (PyTree) – The nested structure of lists or tuples.

  • indices (str | int | Sequence[str | int]) – The indices of the element to return.

  • add (optional) – The element to add to the nested structure of lists or tuples.

  • at (optional) – Specifies the position where to add the element.

Return type:

Any

Returns:

element – The element of the nested structure of lists or tuples.