def render_selection_to_foldable_representation(
    selection: selectors.Selection,
    visible_selection: bool = True,
    ignore_exceptions: bool = False,
) -> rendering_parts.RenderableTreePart:
  """Renders a top-level selection object to its foldable representation.

  This function produces a rendering of either the selection
  itself, or its selected value. However, instead of directly outputting the
  attributes of the selection, we draw the selected regions as highlighted boxes
  around selected subtrees.

  The output is still a round-trippable "repr"-like output. When
  `visible_selection` is True, copying it will reproduce the original selection.
  When `visible_selection` is False, copying it will reproduce the selected
  value.

  Args:
    selection: A selection to render.
    visible_selection: Whether to render the selection itself rather than the
      selected value. If True, renders code to build the selection itself, and
      highlights selected nodes in blue. If False, renders the value but expands
      the selected nodes.
    ignore_exceptions: Whether to catch errors during rendering of subtrees and
      show a fallback for those subtrees, instead of failing the entire
      renderer. Best used in contexts where
      `render_selection_to_foldable_representation` is not being called
      directly, e.g. when registering this as a default pretty-printer.

  Returns:
    A representation of the selection, expanded accordingly.
  """
  # Fetch the default renderer, and add selection-specific handlers.
  base_renderer = treescope.active_renderer.get()
  extended_renderer = base_renderer.extended_with(
      wrapper_hooks=[_wrap_selected_nodes],
  )
  # Enter a scope where we know which values to select, and render with our
  # extended renderer.
  tracker = SelectedNodeTracker(
      selected_node_paths=set(
          jax.tree_util.keystr(keypath)
          for keypath in selection.selected_by_path.keys()
      ),
      rendered_node_paths=set(),
      repeated_node_paths=set(),
      visible_boundary=visible_selection,
  )
  with _selected_nodes.set_scoped(tracker):
    rendered_ir = rendering_parts.build_full_line_with_annotations(
        extended_renderer.to_foldable_representation(
            selection.deselect(),
            ignore_exceptions=ignore_exceptions,
        )
    )

  warnings = []

  unrendered_node_paths = {
      path
      for path in tracker.selected_node_paths
      if path not in tracker.rendered_node_paths
  }
  if unrendered_node_paths:
    warnings.extend([
        "# Some selected nodes were not rendered!",
        "# This likely means the current treescope handlers are not assigning ",
        "# keypaths consistently with JAX for some custom PyTree type.",
        f"# Missing keypaths: {unrendered_node_paths}",
    ])

  if tracker.repeated_node_paths:
    warnings.extend([
        "# Rendered a single selected node multiple times!",
        "# This likely means the current treescope handlers are not assigning ",
        "# keypaths consistently with JAX for some custom PyTree type.",
        f"# Repeated keypaths: {tracker.repeated_node_paths}",
    ])

  # Expand the selected nodes.
  layout_algorithms.expand_to_layout_marks(
      rendered_ir,
      marks=[SelectionBoundaryLayoutMark()],
      collapse_weak_others=True,
  )

  # For convenience, make a default balanced layout for the remaining nodes.
  layout_algorithms.expand_for_balanced_layout(rendered_ir)

  if visible_selection:
    # Render the keypaths:
    keypath_rendering = rendering_parts.build_full_line_with_annotations(
        base_renderer.to_foldable_representation(
            tuple(key for key in selection.selected_by_path.keys()),
            ignore_exceptions=ignore_exceptions,
            root_keypath=None,
        )
    )
    layout_algorithms.expand_for_balanced_layout(keypath_rendering)

    # Combine everything into a rendering of the selection itself.
    count = len(selection)
    result = rendering_parts.siblings(
        rendering_parts.comment_color(rendering_parts.text("pz.select(")),
        rendering_parts.indented_children([rendered_ir]),
        rendering_parts.comment_color(rendering_parts.text(").at_keypaths(")),
        rendering_parts.build_custom_foldable_tree_node(
            label=rendering_parts.comment_color(
                rendering_parts.fold_condition(
                    collapsed=rendering_parts.text(
                        f"<{count} subtrees, highlighted above>"
                    ),
                    expanded=rendering_parts.text(
                        f"# {count} subtrees, highlighted above"
                    ),
                )
            ),
            contents=rendering_parts.fold_condition(
                expanded=rendering_parts.indented_children([keypath_rendering])
            ),
            expand_state=rendering_parts.ExpandState.COLLAPSED,
        ).renderable,
        rendering_parts.comment_color(rendering_parts.text(")")),
    )
  else:
    # Just return our existing rendering.
    result = rendered_ir

  if warnings:
    result = rendering_parts.on_separate_lines([
        rendering_parts.error_color(
            rendering_parts.on_separate_lines(warnings)
        ),
        result,
    ])

  return result
Renders a top-level selection object to its foldable representation.

  This function produces a rendering of either the selection
  itself, or its selected value. However, instead of directly outputting the
  attributes of the selection, we draw the selected regions as highlighted boxes
  around selected subtrees.

  The output is still a round-trippable "repr"-like output. When
  `visible_selection` is True, copying it will reproduce the original selection.
  When `visible_selection` is False, copying it will reproduce the selected
  value.

  Args:
    selection: A selection to render.
    visible_selection: Whether to render the selection itself rather than the
      selected value. If True, renders code to build the selection itself, and
      highlights selected nodes in blue. If False, renders the value but expands
      the selected nodes.
    ignore_exceptions: Whether to catch errors during rendering of subtrees and
      show a fallback for those subtrees, instead of failing the entire
      renderer. Best used in contexts where
      `render_selection_to_foldable_representation` is not being called
      directly, e.g. when registering this as a default pretty-printer.

  Returns:
    A representation of the selection, expanded accordingly.
  
penzai/core/_treescope_handlers/selection_rendering.py