feat: add highlight selection to S2 TreeView

packages/@react-spectrum/s2/chromatic/TreeView.stories.tsx

@@ -143,14 +143,23 @@ export const TreeStatic: StoryObj<typeof TreeExample> = {
   render: (args) => <TreeExample {...args} />
 };
 
-export const TreeSelection: StoryObj<typeof TreeExample> = {
+export const TreeCheckboxSelection: StoryObj<typeof TreeExample> = {
   ...TreeStatic,
   args: {
     selectionMode: 'multiple',
     defaultSelectedKeys: ['projects-2', 'projects-3']
   }
 };
 
+export const TreeHighlightSelection: StoryObj<typeof TreeExample> = {
+  ...TreeStatic,
+  args: {
+    selectionMode: 'multiple',
+    selectionStyle: 'highlight',
+    defaultSelectedKeys: ['projects-2', 'projects-3']
+  }
+};
+
 export const TreeDisableSelection: StoryObj<typeof TreeExample> = {
   ...TreeStatic,
   args: {

packages/@react-spectrum/s2/src/TreeView.tsx

@@ -26,6 +26,7 @@ import {
   TreeItemContentProps,
   TreeLoadMoreItem,
   TreeLoadMoreItemProps,
+  TreeState,
   useContextProps,
   Virtualizer
 } from 'react-aria-components';
@@ -39,7 +40,7 @@ import {IconContext} from './Icon';
 import intlMessages from '../intl/*.json';
 import {ProgressCircle} from './ProgressCircle';
 import {raw} from '../style/style-macro' with {type: 'macro'};
-import React, {createContext, forwardRef, JSXElementConstructor, ReactElement, ReactNode, useRef} from 'react';
+import React, {createContext, forwardRef, JSXElementConstructor, ReactElement, ReactNode, useContext, useRef} from 'react';
 import {Text, TextContext} from './Content';
 import {useActionBarContainer} from './ActionBar';
 import {useDOMRef} from '@react-spectrum/utils';
@@ -53,7 +54,15 @@ interface S2TreeProps {
   renderActionBar?: (selectedKeys: 'all' | Set<Key>) => ReactElement
 }
 
-export interface TreeViewProps<T> extends Omit<RACTreeProps<T>, 'style' | 'className' | 'render' | 'onRowAction' | 'selectionBehavior' | 'onScroll' | 'onCellAction' | 'dragAndDropHooks' | keyof GlobalDOMAttributes>, UnsafeStyles, S2TreeProps {
+interface TreeViewStyleProps {
+  /**
+   * How selection should be displayed.
+   * @default 'checkbox'
+   */
+  selectionStyle?: 'highlight' | 'checkbox'
+}
+
+export interface TreeViewProps<T> extends Omit<RACTreeProps<T>, 'style' | 'className' | 'render' | 'onRowAction' | 'selectionBehavior' | 'onScroll' | 'onCellAction' | 'dragAndDropHooks' | keyof GlobalDOMAttributes>, UnsafeStyles, S2TreeProps, TreeViewStyleProps {
   /** Spectrum-defined styles, returned by the `style()` macro. */
   styles?: StylesPropWithHeight
 }
@@ -109,11 +118,13 @@ const tree = style({
   }
 });
 
+let InternalTreeViewContext = createContext<{selectionStyle?: 'highlight' | 'checkbox'}>({});
+
 /**
  * A tree view provides users with a way to navigate nested hierarchical information.
  */
 export const TreeView = /*#__PURE__*/ (forwardRef as forwardRefType)(function TreeView<T extends object>(props: TreeViewProps<T>, ref: DOMRef<HTMLDivElement>) {
-  let {children, UNSAFE_className, UNSAFE_style} = props;
+  let {children, selectionStyle = 'checkbox', UNSAFE_className, UNSAFE_style} = props;
   let scale = useScale();
 
   let renderer;
@@ -137,20 +148,22 @@ export const TreeView = /*#__PURE__*/ (forwardRef as forwardRefType)(function Tr
           rowHeight: scale === 'large' ? 50 : 40
         }}>
         <TreeRendererContext.Provider value={{renderer}}>
-          <Tree
-            {...props}
-            style={{
-              paddingBottom: actionBarHeight > 0 ? actionBarHeight + 8 : 0,
-              scrollPaddingBottom: actionBarHeight > 0 ? actionBarHeight + 8 : 0
-            }}
-            className={tree}
-            selectionBehavior="toggle"
-            selectedKeys={selectedKeys}
-            defaultSelectedKeys={undefined}
-            onSelectionChange={onSelectionChange}
-            ref={scrollRef}>
-            {props.children}
-          </Tree>
+          <InternalTreeViewContext.Provider value={{selectionStyle}}>
+            <Tree
+              {...props}
+              style={{
+                paddingBottom: actionBarHeight > 0 ? actionBarHeight + 8 : 0,
+                scrollPaddingBottom: actionBarHeight > 0 ? actionBarHeight + 8 : 0
+              }}
+              className={tree}
+              selectionBehavior={selectionStyle === 'highlight' ? 'replace' : 'toggle'}
+              selectedKeys={selectedKeys}
+              defaultSelectedKeys={undefined}
+              onSelectionChange={onSelectionChange}
+              ref={scrollRef}>
+              {props.children}
+            </Tree>
+          </InternalTreeViewContext.Provider>
         </TreeRendererContext.Provider>
       </Virtualizer>
       {actionBar}
@@ -162,7 +175,7 @@ const rowBackgroundColor = {
   default: '--s2-container-bg',
   isFocusVisibleWithin: colorMix('gray-25', 'gray-900', 7),
   isHovered: colorMix('gray-25', 'gray-900', 7),
-  isPressed: colorMix('gray-25', 'gray-900', 10),
+  isPressed: colorMix('gray-25', 'gray-900', 7),
   isSelected: {
     default: colorMix('gray-25', 'gray-900', 7),
     isFocusVisibleWithin: colorMix('gray-25', 'gray-900', 10),
@@ -175,6 +188,17 @@ const rowBackgroundColor = {
 } as const;
 
 const treeRow = style({
+  ...focusRing(),
+  outlineOffset: -2,
+  outlineColor: {
+    default: 'focus-ring',
+    forcedColors: 'Highlight',
+    selectionStyle: {
+      highlight: {
+        forcedColors: 'ButtonBorder'
+      }
+    }
+  },
   position: 'relative',
   display: 'flex',
   height: 40,
@@ -186,7 +210,6 @@ const treeRow = style({
     isSelected: baseColor('neutral'),
     forcedColors: 'ButtonText'
   },
-  outlineStyle: 'none',
   cursor: {
     default: 'default',
     isLink: 'pointer'
@@ -195,33 +218,39 @@ const treeRow = style({
     type: 'backgroundColor',
     value: rowBackgroundColor
   },
-  '--rowFocusIndicatorColor': {
-    type: 'outlineColor',
-    value: {
-      default: 'focus-ring',
-      forcedColors: 'Highlight'
-    }
-  }
+  '--borderRadiusTreeItem': {
+    type: 'borderTopStartRadius',
+    value: 'sm'
+  },
+  borderRadius: 'sm'
 });
 
 const treeCellGrid = style({
   display: 'grid',
   width: 'full',
   height: 'full',
   boxSizing: 'border-box',
+  borderRadius: 'sm',
   alignContent: 'center',
   alignItems: 'center',
   gridTemplateColumns: ['auto', 'auto', 'auto', 'auto', 'auto', '1fr', 'minmax(0, auto)', 'auto'],
   gridTemplateRows: '1fr',
   gridTemplateAreas: [
     'drag-handle checkbox level-padding expand-button icon content actions actionmenu'
   ],
-  backgroundColor: '--rowBackgroundColor',
   paddingEnd: 4, // account for any focus rings on the last item in the cell
   color: {
     isDisabled: {
       default: 'gray-400',
       forcedColors: 'GrayText'
+    },
+    forcedColors: 'ButtonText',
+    selectionStyle: {
+      highlight: {
+        isSelected: {
+          forcedColors: 'HighlightText'
+        }
+      }
     }
   },
   '--rowSelectedBorderColor': {
@@ -238,6 +267,82 @@ const treeCellGrid = style({
       default: 'focus-ring',
       forcedColors: 'Highlight'
     }
+  },
+  '--borderColor': {
+    type: 'borderColor',
+    value: {
+      default: 'transparent',
+      selectionStyle: {
+        highlight: 'blue-900'
+      },
+      forcedColors: 'ButtonBorder'
+    }
+  },
+  forcedColorAdjust: 'none'
+});
+
+const treeRowBackground = style({
+  position: 'absolute',
+  zIndex: -1,
+  inset: 0,
+  backgroundColor: {
+    default: '--rowBackgroundColor',
+    forcedColors: 'Background',
+    selectionStyle: {
+      highlight: {
+        default: '--rowBackgroundColor',
+        isSelected: {
+          default: colorMix('gray-25', 'blue-900', 10),
+          isHovered: colorMix('gray-25', 'blue-900', 15),
+          isPressed: colorMix('gray-25', 'blue-900', 15),
+          isFocusVisible: colorMix('gray-25', 'blue-900', 15),
+          forcedColors: 'Highlight'
+        }
+      }
+    }
+  },
+  borderTopStartRadius: {
+    default: '--borderRadiusTreeItem',
+    isPreviousSelected: {
+      default: '--borderRadiusTreeItem',
+      isSelected: 'none'
+    }
+  },
+  borderTopEndRadius: {
+    default: '--borderRadiusTreeItem',
+    isPreviousSelected: {
+      default: '--borderRadiusTreeItem',
+      isSelected: 'none'
+    }
+  },
+  borderBottomStartRadius: {
+    default: '--borderRadiusTreeItem',
+    isNextSelected: {
+      default: '--borderRadiusTreeItem',
+      isSelected: 'none'
+    }
+  },
+  borderBottomEndRadius: {
+    default: '--borderRadiusTreeItem',
+    isNextSelected: {
+      default: '--borderRadiusTreeItem',
+      isSelected: 'none'
+    }
+  },
+  borderTopWidth: {
+    default: 1,
+    isPreviousSelected: 0
+  },
+  borderBottomWidth: {
+    default: 1,
+    isNextSelected: 0
+  },
+  borderStartWidth: 1,
+  borderEndWidth: 1,
+  borderStyle: 'solid',
+  borderColor: {
+    default: 'transparent',
+    isSelected: '--borderColor'
   }
 });
 
@@ -298,13 +403,16 @@ export const TreeViewItem = (props: TreeViewItemProps): ReactNode => {
   let {
     href
   } = props;
+  let {selectionStyle} = useContext(InternalTreeViewContext);
 
   return (
     <TreeItem
       {...props}
       className={(renderProps) => treeRow({
         ...renderProps,
-        isLink: !!href
+        isLink: !!href,
+        selectionStyle,
+        isPreviousSelected: isPrevSelected(renderProps.id, renderProps.state)
       }) + (renderProps.isFocusVisible ? ' ' + treeRowFocusIndicator : '')} />
   );
 };
@@ -320,19 +428,15 @@ export const TreeViewItemContent = (props: TreeViewItemContentProps): ReactNode
   } = props;
   let scale = useScale();
 
+  let {selectionStyle} = useContext(InternalTreeViewContext);
+
   return (
     <TreeItemContent>
-      {({isExpanded, hasChildItems, selectionMode, selectionBehavior, isDisabled, isSelected, id, state}) => {
-        let isNextSelected = false;
-        let isNextFocused = false;
-        let keyAfter = state.collection.getKeyAfter(id);
-        if (keyAfter != null) {
-          isNextSelected = state.selectionManager.isSelected(keyAfter);
-        }
-        let isFirst = state.collection.getFirstKey() === id;
+      {({isExpanded, hasChildItems, selectionMode, selectionBehavior, isDisabled, isSelected, id, state, isHovered, isFocusVisible}) => {
         return (
-          <div className={treeCellGrid({isDisabled, isNextSelected, isSelected, isFirst, isNextFocused})}>
-            {selectionMode !== 'none' && selectionBehavior === 'toggle' && (
+          <div className={treeCellGrid({isDisabled, isNextSelected: isNextSelected(id, state), isSelected, selectionStyle})}>
+            <div className={treeRowBackground({isHovered, isFocusVisible, isSelected, selectionStyle, isNextSelected: isNextSelected(id, state), isPreviousSelected: isPrevSelected(id, state)})} />
+            {selectionMode !== 'none' && selectionBehavior === 'toggle' && selectionStyle !== 'highlight' && (
               // TODO: add transition?
               <div className={treeCheckbox({isDisabled: isDisabled || !state.selectionManager.canSelectItem(id) || state.disabledKeys.has(id)})}>
                 <Checkbox slot="selection" />
@@ -459,3 +563,27 @@ function ExpandableRowChevron(props: ExpandableRowChevronProps) {
     </Button>
   );
 }
+
+function isNextSelected(id: Key | undefined, state: TreeState<unknown>) {
+  if (id == null || !state) {
+    return false;
+  }
+  let keyAfter = state.collection.getKeyAfter(id);
+
+  // We need to skip non-item nodes because the selection manager will map non-item nodes to their parent before checking selection
+  let node = keyAfter ? state.collection.getItem(keyAfter) : null;
+  while (node && node.type !== 'item' && keyAfter) {
+    keyAfter = state.collection.getKeyAfter(keyAfter);
+    node = keyAfter ? state.collection.getItem(keyAfter) : null;
+  }
+
+  return keyAfter != null && state.selectionManager.isSelected(keyAfter);
+}
+
+function isPrevSelected(id: Key | undefined, state: TreeState<unknown>) {
+  if (id == null || !state) {
+    return false;
+  }
+  let keyBefore = state.collection.getKeyBefore(id);
+  return keyBefore != null && state.selectionManager.isSelected(keyBefore);
+}

packages/dev/s2-docs/pages/s2/TreeView.mdx

@@ -14,7 +14,7 @@ export const description = 'Displays hierarchical data with selection and collap
 
 <PageDescription>{docs.exports.TreeView.description}</PageDescription>
 
-```tsx render docs={docs.exports.TreeView} links={docs.links} props={['selectionMode']} initialProps={{selectionMode: 'multiple'}} type="s2" wide
+```tsx render docs={docs.exports.TreeView} links={docs.links} props={['selectionMode', 'selectionStyle']} initialProps={{selectionMode: 'multiple', selectionStyle: 'checkbox'}} type="s2" wide
 "use client";
 import {TreeView, TreeViewItem, TreeViewItemContent} from '@react-spectrum/s2';
 import {style} from '@react-spectrum/s2/style' with {type: 'macro'};