Module utilities.tree_utils

Tree manipulation utilities with no external dependencies

This module provides methods for building and manipulating trees.

A Tree is a nested dictionary. A Leaf is any other object. A Node is a non-leaf Tree node. A TreeDef is a nested dictionary with no data, only structure.

Functions overview

branches(tree[, prefix])

Generate all branches (paths from root node to any leaf) in a tree

get_nested(tree, branch)

Get a value from a tree branch, specifying a branch

make_prototype_tree(full_tree, target_tree)

Construct a tree with boolean leaves, for nodes that match a target tree

set_matching(full_tree, target_tree, value)

Set the values in a full tree, for branches that match a target tree

set_matching_select(full_tree, target_tree, ...)

Set the values in a full tree, for branches that match a target tree, if the target tree leaf nodes evaluate to True

set_nested(tree, branch, value[, inplace])

Set a value in a tree branch, specifying a branch

tree_find(tree)

Generate the tree branches to tree nodes that evaluate to True

tree_flatten(tree[, leaves])

Flatten a tree into a linear list of leaf nodes and a tree description

tree_map(tree, f)

Map a function over the leaves of a tree

tree_unflatten(treedef, leaves[, leaves_tail])

Build a tree from a flat list of leaves, plus a tree definition

tree_update(target, additional[, inplace])

Perform a recursive update of a tree to insert or replace nodes from a second tree

Functions

utilities.tree_utils.branches(tree: Iterable | MutableMapping | Mapping, prefix: list | None = None) Generator[Tuple, None, None][source]

Generate all branches (paths from root node to any leaf) in a tree

Parameters:
  • tree (Tree) – A nested dictionary tree structure

  • prefix (list) – The current branch prefix. Default: None (this is the root)

Yields:

Tuple[str]

utilities.tree_utils.get_nested(tree: Iterable | MutableMapping | Mapping, branch: Tuple) None[source]

Get a value from a tree branch, specifying a branch

Parameters:
  • tree (Tree) – A nested dictionary tree structure

  • branch (Tuple[str]) – A branch: a tuple of indices to walk through the tree

utilities.tree_utils.make_prototype_tree(full_tree: Iterable | MutableMapping | Mapping, target_tree: Iterable | MutableMapping | Mapping) Iterable | MutableMapping | Mapping[source]

Construct a tree with boolean leaves, for nodes that match a target tree

Make a prototype tree, indicating which nodes in a large tree should be selected for analysis or processing. This is done on the basis of a smaller “target” tree, which contains only the leaf nodes of interest.

Examples

>>> target_tree = {'a': 0, 'b': {'b2': 0}}
>>> full_tree = {'a': 1, 'b': {'b1': 2, 'b2': 3}, 'c': 4, 'd': 5}
>>> make_prototype_tree(full_tree, target_tree)
{'a': True, 'b': {'b1': False, 'b2': True}, 'c': False, 'd': False}
Parameters:
  • full_tree (Tree) – A large tree to search through.

  • target_tree (Tree) – A tree with only few leaf nodes. These nodes will be identifed within the full tree.

Returns:

A nested tree with the same tree structure as full_tree, but with bool leaf nodes. Leaf nodes will be True for branches matching those specified in target_tree, and False otherwise.

Return type:

Tree

utilities.tree_utils.set_matching(full_tree: Iterable | MutableMapping | Mapping, target_tree: Iterable | MutableMapping | Mapping, value: Any, inplace: bool = False) Iterable | MutableMapping | Mapping[source]

Set the values in a full tree, for branches that match a target tree

Parameters:
  • full_tree (Tree) – A tree to search over. The values in this tree will be replaced with value

  • target_tree (Tree) – A tree that defines the target branches to set in full_tree. Matching branches in full_tree will have their values replaced with value

  • value (Any) – The value to set in full_tree.

  • inplace (bool) – If False (default), a copy of the tree will be returned. If True, the operation will be performed in place, and the original tree will be returned

Returns:

The modified tree

Return type:

Tree

utilities.tree_utils.set_matching_select(full_tree: Iterable | MutableMapping | Mapping, target_tree: Iterable | MutableMapping | Mapping, value: Any, inplace: bool = False) Iterable | MutableMapping | Mapping[source]

Set the values in a full tree, for branches that match a target tree, if the target tree leaf nodes evaluate to True

Parameters:
  • full_tree (Tree) – A tree to search over. The values in this tree will be replaced with value

  • target_tree (Tree) – A tree that defines the target branches to set in full_tree. Matching branches in full_tree will have their values replaced with value, if the leaf node in target_tree` evaluates to ``True

  • value (Any) – The value to set in full_tree.

  • inplace (bool) – If False (default), a copy of the tree will be returned. If True, the operation will be performed in place, and the original tree will be returned

Returns:

The modified tree

Return type:

Tree

utilities.tree_utils.set_nested(tree: Iterable | MutableMapping | Mapping, branch: Tuple, value: Any, inplace: bool = False) Iterable | MutableMapping | Mapping[source]

Set a value in a tree branch, specifying a branch

The leaf node must already exist in the tree.

Parameters:
  • tree (Tree) – A nested dictionary tree structure

  • branch (Tuple[str]) – A branch: a tuple of indices to walk through the tree

  • value (Any) – The value to set at the tree leaf

  • inplace (bool) – If False (default), a copy of the tree will be returned. If True, the operation will be performed in place, and the original tree will be returned

Returns:

The modified tree

Return type:

Tree

utilities.tree_utils.tree_find(tree: Iterable | MutableMapping | Mapping) Generator[Tuple, None, None][source]

Generate the tree branches to tree nodes that evaluate to True

Parameters:

tree (Tree) – A tree to examine

Returns:

A list of all tree branches, for which the corresponding tree leaf evaluate to True

Return type:

list

utilities.tree_utils.tree_flatten(tree: Iterable | MutableMapping | Mapping, leaves: List[Any] | None = None) Tuple[List[Any], Dict][source]

Flatten a tree into a linear list of leaf nodes and a tree description

This function operates similar to jax.tree_utils.tree_flatten, but is not directly compatible.

A Tree tree will be serialised into a simple list of leaf nodes, which can then be conveniently manipulated. A TreeDef will also be returned, which is a nested dictionary with the same structure as tree.

The function tree_unflatten() performs the reverse operation.

Parameters:
  • tree (Tree) – A tree to flatten

  • leaves (Optional[List[Any]]) – Used recursively. Should be left as None by the user.

Returns:

A list of leaf nodes from the flattened tree, and a tree definition.

Return type:

Tuple[List[Any], TreeDef]

utilities.tree_utils.tree_map(tree: Iterable | MutableMapping | Mapping, f: Callable) Iterable | MutableMapping | Mapping[source]

Map a function over the leaves of a tree

This function performs a recurdive depth-first traversal of the tree.

Parameters:
  • tree (Tree) – A tree to traverse

  • f (Callable) – A function which is called on each leaf of the tree. Must have the signature Callable[Leaf] -> Any

Returns:

A tree with the same structure as tree, with leaf nodes replaced with the result of calling f on each leaf.

Return type:

Tree

utilities.tree_utils.tree_unflatten(treedef: Dict, leaves: List, leaves_tail: List[Any] | None = None) Iterable | MutableMapping | Mapping[source]

Build a tree from a flat list of leaves, plus a tree definition

This function takes a flattened tree representation, as built by tree_flatten(), and reconstructs a matching Tree structure.

Parameters:
  • treedef (TreeDef) – A tree definition as returned by tree_flatten()

  • leaves (List[Any]) – A list of leaf nodes to use in constructing the tree

  • leaves_tail (Optional[List[Any]]) – Used recursively. Should be left as None by the end user

Returns:

The reconstructed tree, with leaves taken from leaves

Return type:

Tree

utilities.tree_utils.tree_update(target: Iterable | MutableMapping | Mapping, additional: Iterable | MutableMapping | Mapping, inplace: bool = False) Iterable | MutableMapping | Mapping[source]

Perform a recursive update of a tree to insert or replace nodes from a second tree

Requires a target Tree and a source Tree additional, which will provide the source data to update in target.

Both target and additional will be traversed depth-first simultaneously. Leaf nodes that exist in target but not in additional will not be modified. Leaf nodes that exist in additional but not in target will be inserted into target at the corresponding location. Leaf nodes that exist in both trees will have their data updated from additional to target, using the python update() function.

Parameters:
  • target (Tree) – The tree to update.

  • additional (Tree) – The source tree to insert / replace nodes from, into target. Will not be modified.

  • inplace (bool) – If False (default), a copy of the tree will be returned. If True, the operation will be performed in place, and the original tree will be returned.

Returns:

The modified target tree

Return type:

Tree