def apply_with_selected_index(
      self,
      fn: Callable[[int, SelectedSubtree], Any],
      keep_selected: bool = False,
  ) -> Any | Selection:
    """Applies a function, passing both the selected nodes and their indices.

    Indices are taken relative to the linearized sequence of currently selected
    nodes. In other words, if there are five nodes selected, then ``fn`` will
    be called with the numbers 0 through 4 inclusive, regardless of the specific
    keypaths to the nodes.

    Args:
      fn: Function to call. This function will be passed both the index of the
        selected node and the value, and should return something to replace the
        value with.
      keep_selected: Whether to keep the node selected after the transformation.

    Returns:
      The tree or modified selection after inserting the results of ``fn``,
      depending on ``keep_selected``.
    """
    new_selected_by_path = collections.OrderedDict({
        key: fn(i, value)
        for i, (key, value) in enumerate(self.selected_by_path.items())
    })
    new_selection = Selection(
        selected_by_path=new_selected_by_path, remainder=self.remainder
    )
    if keep_selected:
      return new_selection
    else:
      return new_selection.deselect()
Applies a function, passing both the selected nodes and their indices.

    Indices are taken relative to the linearized sequence of currently selected
    nodes. In other words, if there are five nodes selected, then ``fn`` will
    be called with the numbers 0 through 4 inclusive, regardless of the specific
    keypaths to the nodes.

    Args:
      fn: Function to call. This function will be passed both the index of the
        selected node and the value, and should return something to replace the
        value with.
      keep_selected: Whether to keep the node selected after the transformation.

    Returns:
      The tree or modified selection after inserting the results of ``fn``,
      depending on ``keep_selected``.
    
penzai/core/selectors.py