Method flatten_selected_selections

penzai/core/selectors.py:435–474

Flattens a selection whose selected values are all selections. This function takes a selection for which all of the selected values are already selections, and merges them into a single selection that selects all of the values from each individual selection. You can use this to bui

(
      self: "Selection[Selection[SelectedSubtree]]",
  )

    return self.set_by_path(replacer)

  def flatten_selected_selections(
      self: "Selection[Selection[SelectedSubtree]]",
  ) -> "Selection[SelectedSubtree]":
    """Flattens a selection whose selected values are all selections.

    This function takes a selection for which all of the selected values are
    already selections, and merges them into a single selection that selects
    all of the values from each individual selection.

    You can use this to build more complex selections by chaining your own
    logic. For instance, if you have written a function ``f`` that selects part
    of a tree, you can run

    ::

      selection.apply(f, keep_selected=True).flatten_selected_selections()

    to "broadcast" that logic across all of the already-selected subtrees in
    the original selection.

    See also `refine`, which allows you to express similar transformations more
    directly.

    Returns:
      A flattened selection object.
    """
    # Strategy: Replace the values of each inner selection by a boundary,
    # deselect everything, then re-select at the boundary.

    def process_subselection(subselection: Selection) -> Any:
      if not isinstance(subselection, Selection):
        raise ValueError(
            "flatten_selected_selections can only be called on Selections for"
            " which all values in `.selected_by_path` are also Selections. Got"
            f" {subselection}"
        )
      return subselection.apply(_InProgressSelectionBoundary)

    with _wrap_selection_errors(self):
      return _build_selection_from_boundary(self.apply(process_subselection))

  def refine(
      self, selector_fn: "Callable[[Any], Selection[OtherSubtree]]"