def _wrap_selected_nodes(
    node: Any,
    path: str | None,
    node_renderer: renderers.TreescopeSubtreeRenderer,
) -> (
    rendering_parts.RenderableTreePart
    | rendering_parts.RenderableAndLineAnnotations
    | type(NotImplemented)
):
  """Custom wrapper hook that intercepts selected nodes."""
  tracker = _selected_nodes.get()

  assert tracker is not None
  if path is not None and path in tracker.selected_node_paths:
    if path in tracker.rendered_node_paths:
      tracker.repeated_node_paths.add(path)
    else:
      tracker.rendered_node_paths.add(path)

    # Tag the child, and possibly annotate its visualization.
    rendering = node_renderer(node, path)

    tagged_rendering = rendering_parts.RenderableAndLineAnnotations(
        renderable=rendering_parts.with_layout_mark(
            rendering.renderable, SelectionBoundaryLayoutMark()
        ),
        annotations=rendering.annotations,
    )

    if tracker.visible_boundary:
      wrapped_rendering = rendering_parts.siblings_with_annotations(
          rendering_parts.build_custom_foldable_tree_node(
              contents=rendering_parts.in_outlined_box(
                  rendering_parts.on_separate_lines([
                      rendering_parts.fold_condition(
                          expanded=rendering_parts.custom_style(
                              rendering_parts.text("# Selected:"),
                              css_style="color: darkcyan; font-weight: bold;",
                          )
                      ),
                      tagged_rendering.renderable,
                  ]),
                  css_style=(
                      "border: 2px solid cyan; border-left: 2ch solid cyan;"
                  ),
              )
          ),
          extra_annotations=[tagged_rendering.annotations],
      )
    else:
      wrapped_rendering = tagged_rendering

    return wrapped_rendering

  return NotImplemented
Custom wrapper hook that intercepts selected nodes.
penzai/core/_treescope_handlers/selection_rendering.py