Skip to content

treex.Map

Mixin that adds a .map() method to the class.

map(self, f, *filters, *, inplace=False, flatten_mode=None, is_leaf=None)

map is a wrapper over treeo.map that passes self as the second argument.

Parameters:

Name Type Description Default
f Callable

The function to apply to the leaves.

required
*filters Union[Type[Any], Callable[[FieldInfo], bool]]

The filters used to select the leaves to which the function will be applied.

()
inplace bool

If True, the input obj is mutated and returned.

False
flatten_mode Union[treeo.tree.FlattenMode, str]

Sets a new FlattenMode context for the operation, if None the current context is used.

None

Returns:

Type Description
~A

A new pytree with the changes applied. If inplace is True, the input obj is returned.

Source code in treeo/mixins.py
def map(
    self: A,
    f: tp.Callable,
    *filters: api.Filter,
    inplace: bool = False,
    flatten_mode: tp.Union[api.FlattenMode, str, None] = None,
    is_leaf: tp.Callable[[tp.Any], bool] = None,
) -> A:
    """
    `map` is a wrapper over `treeo.map` that passes `self` as the second argument.

    Arguments:
        f: The function to apply to the leaves.
        *filters: The filters used to select the leaves to which the function will be applied.
        inplace: If `True`, the input `obj` is mutated and returned.
        flatten_mode: Sets a new `FlattenMode` context for the operation, if `None` the current context is used.

    Returns:
        A new pytree with the changes applied. If `inplace` is `True`, the input `obj` is returned.
    """
    return api.map(
        f,
        self,
        *filters,
        inplace=inplace,
        flatten_mode=flatten_mode,
        is_leaf=is_leaf,
    )