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

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

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

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

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