Skip to main content

TreeSpec

The structure of a pytree without its leaves.

A spec comes from :func:~clika_runtime.pytree.tree_flatten or :func:~clika_runtime.pytree.tree_structure, from the treespec_* constructors, or from this constructor in the (type, context, children_specs) form::

>>> from clika_runtime import pytree
>>> leaves, spec = pytree.tree_flatten({"a": [1, (2, None)], "b": 3})
>>> leaves
[1, 2, None, 3]
>>> spec
TreeSpec({'a': [*, (*, *)], 'b': *})
>>> spec.num_leaves, spec.num_nodes, spec.num_children
(4, 7, 2)
>>> spec.unflatten([10, 20, 30, 40])
{'a': [10, (20, 30)], 'b': 40}
>>> pytree.TreeSpec(dict, ["a", "b"], [pytree.LeafSpec(), pytree.LeafSpec()])
TreeSpec({'a': *, 'b': *})

none_is_leaf (default True) records whether None was a leaf when the tree was flattened; with False a None is a node with no children and contributes no leaf. Two specs compare equal only with the same none_is_leaf. namespace is the registry namespace the tree was flattened under; the global namespace ("") is compatible with every other one.

children_specs (property)​

The specs of the root's children, in order (a list copy).

context (property)​

The root container's context: the keys of a dictionary, the class of a namedtuple, the maxlen of a deque, a custom node's own context; None for a leaf, a tuple or a list.

kind (property)​

The root node's :class:~clika_runtime.pytree.TreeKind.

namespace (property)​

The registry namespace the tree was flattened under.

none_is_leaf (property)​

Whether None was treated as a leaf when this spec was built.

num_children (property)​

How many children the root container has.

num_leaves (property)​

How many leaves a tree of this structure has.

num_nodes (property)​

How many nodes the structure has, leaves and containers together.

type (property)​

The root container's type; None for a leaf spec.

__init__​

__init__(self, type: 'type | None' = None, context: 'Any' = None, children_specs: 'Iterable[TreeSpec]' = (), *, none_is_leaf: 'bool' = True, namespace: 'str' = '') -> 'None'

Initialize self. See help(type(self)) for accurate signature.

accessors​

accessors(self) -> 'list[Any]'

One :class:~clika_runtime.pytree.PyTreeAccessor per leaf, in leaf order: a callable path that fetches the leaf from a tree of this structure.

broadcast_to_common_suffix​

broadcast_to_common_suffix(self, other: 'TreeSpec') -> 'TreeSpec'

The smallest spec both this one and other are prefixes of: at every position, the deeper of the two structures.

a = pytree.tree_structure([1, (2, 3), 4]) b = pytree.tree_structure([5, 6, (7, 8)]) a.broadcast_to_common_suffix(b) TreeSpec([, (, ), (, *)])

Raises: ValueError: Where the two structures disagree (different node types, arities, keys, or contexts at the same position).

child​

child(self, index: 'int') -> 'TreeSpec'

The spec of the root's index-th child (negative indices count from the end).

children​

children(self) -> 'list[TreeSpec]'

The specs of the root's children, in order.

compose​

compose(self, inner: 'TreeSpec') -> 'TreeSpec'

The spec of a tree with this structure whose every leaf is a tree of inner's structure.

pytree.tree_structure([1, 2]).compose(pytree.tree_structure((1, 2))) TreeSpec([(*, ), (, *)])

entries​

entries(self) -> 'list[Any]'

The root's child entries: indices for sequences, keys for dictionaries, the names a custom node's flatten function supplied.

entry​

entry(self, index: 'int') -> 'Any'

The entry of the root's index-th child.

flatten_up_to​

flatten_up_to(self, tree: 'Any') -> 'list[Any]'

Flatten tree only as deep as this spec goes, and return the subtrees standing where this spec has leaves.

tree must have this spec as a prefix: the same containers, arities and keys down to this spec's leaves; anything below a leaf is returned whole. This is how :func:~clika_runtime.pytree.tree_map pairs the leaves of its first tree with the matching subtrees of the rest.

spec = pytree.tree_structure({"a": 1, "b": 2}) spec.flatten_up_to({"a": [1, 2], "b": {"c": 3}}) [[1, 2], {'c': 3}]

Raises: ValueError: At the first node whose type, arity, keys, or context differs from this spec; the message names that node.

is_leaf​

is_leaf(self, *, strict: 'bool' = True) -> 'bool'

Whether this spec is a single leaf. With strict=False a childless None node counts too.

is_one_level​

is_one_level(self) -> 'bool'

Whether the root's children are all leaves (a leaf spec itself is not one-level).

is_prefix​

is_prefix(self, other: 'TreeSpec', *, strict: 'bool' = False) -> 'bool'

Whether other can be built by replacing some of this spec's leaves with subtrees.

This is the relation multi-tree :func:~clika_runtime.pytree.tree_map needs between its first tree and the rest. Dictionary key order and a deque maxlen do not matter here (they do matter to ==). With strict=True the two specs must also differ.

pytree.tree_structure([1, 2]).is_prefix(pytree.tree_structure([1, (2, 3)])) True

is_suffix​

is_suffix(self, other: 'TreeSpec', *, strict: 'bool' = False) -> 'bool'

Whether this spec can be built by replacing some of other's leaves with subtrees; the reverse of :meth:is_prefix.

keypaths​

keypaths(self) -> 'list[tuple[Any, ...]]'

One key path per leaf, in leaf order (see :func:~clika_runtime.pytree.keystr).

one_level​

one_level(self) -> 'TreeSpec | None'

This spec's root with every child replaced by a leaf; None for a leaf spec.

paths​

paths(self) -> 'list[tuple[Any, ...]]'

One tuple of entries per leaf, root to leaf, in leaf order.

transform​

transform(self, f_node: 'Callable[[TreeSpec], TreeSpec] | None' = None, f_leaf: 'Callable[[TreeSpec], TreeSpec] | None' = None) -> 'TreeSpec'

Rebuild the spec node by node.

f_node receives each container as a one-level spec (its children replaced by leaves) and returns a one-level spec with the same number of children, which takes its place; f_leaf receives each leaf spec and returns any spec, which is spliced in. None leaves that side unchanged.

spec = pytree.tree_structure({"a": (1, 2), "b": 3}) spec.transform(lambda s: pytree.treespec_list(s.children())) TreeSpec([[*, *], ]) spec.transform(f_leaf=lambda s: pytree.tree_structure([1, 2])) TreeSpec({'a': ([, ], [, ]), 'b': [, *]})

Raises: TypeError: If a function returns something other than a TreeSpec. ValueError: If a returned spec has a different none_is_leaf, an incompatible namespace, or (for f_node) is not one-level with the same number of children.

traverse​

traverse(self, leaves: 'Iterable[Any]', f_node: 'Callable[[Any], Any] | None' = None, f_leaf: 'Callable[[Any], Any] | None' = None) -> 'Any'

Like :meth:walk, but f_node receives each REBUILT container (children already placed) and returns its replacement.

unflatten​

unflatten(self, leaves: 'Iterable[Any]') -> 'Any'

Put leaves back into this structure.

Raises: ValueError: If leaves does not hold exactly :attr:num_leaves values; the message states both counts.

walk​

walk(self, leaves: 'Iterable[Any]', f_node: 'Callable[[builtins.type | None, Any, tuple[Any, ...]], Any] | None' = None, f_leaf: 'Callable[[Any], Any] | None' = None) -> 'Any'

Fold the structure bottom-up over leaves without rebuilding the containers.

f_leaf(leaf) transforms each leaf (identity when None); f_node(node_type, context, children) combines a container's transformed children (children is a tuple) into one value. When f_node is None the containers are rebuilt as in :meth:unflatten.

spec = pytree.tree_structure({"a": (1, 2), "b": 3}) spec.walk([1, 2, 3], lambda t, ctx, children: sum(children), lambda x: x * 10) 60

Raises: ValueError: If leaves does not hold exactly :attr:num_leaves values.