feat: Support keyboard focusing loading spinners in collection components

packages/@react-aria/grid/src/GridKeyboardDelegate.ts

@@ -59,7 +59,7 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
   }
 
   private isDisabled(item: Node<unknown>) {
-    return this.disabledBehavior === 'all' && (item.props?.isDisabled || this.disabledKeys.has(item.key));
+    return this.disabledBehavior === 'all' && (item.props?.isDisabled || this.disabledKeys.has(item.key) || (item.type === 'loader' && !item.props.isLoading));
   }
 
   protected findPreviousKey(fromKey?: Key, pred?: (item: Node<T>) => boolean): Key | null {
@@ -148,10 +148,10 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     }
 
     // Find the next item
-    key = this.findNextKey(key, (item => item.type === 'item'));
+    key = this.findNextKey(key, (item => item.type === 'item' || item.type === 'loader'));
     if (key != null) {
       // If focus was on a cell, focus the cell with the same index in the next row.
-      if (this.isCell(startItem)) {
+      if (this.isCell(startItem) && this.collection.getItem(key)?.type === 'item') {
         let startIndex = startItem.colIndex ? startItem.colIndex : startItem.index;
         return this.getKeyForItemInRowByIndex(key, startIndex);
       }
@@ -334,7 +334,7 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     }
 
     // Find the last row
-    key = this.findPreviousKey(undefined, item => item.type === 'item');
+    key = this.findPreviousKey(undefined, item => item.type === 'item' || item.type === 'loader');
 
     // If global flag is set (or if focus mode is cell), focus the last cell in the last row.
     if (key != null && ((item && this.isCell(item) && global) || this.focusMode === 'cell')) {

packages/@react-aria/grid/src/useGrid.ts

@@ -12,6 +12,7 @@
 
 import {AriaLabelingProps, DOMAttributes, DOMProps, Key, KeyboardDelegate, RefObject} from '@react-types/shared';
 import {filterDOMProps, mergeProps, useId} from '@react-aria/utils';
+import {getItemCount} from '@react-stately/collections';
 import {GridCollection} from '@react-types/grid';
 import {GridKeyboardDelegate} from './GridKeyboardDelegate';
 import {gridMap} from './utils';
@@ -175,7 +176,7 @@ export function useGrid<T>(props: GridProps, state: GridState<T, GridCollection<
   );
 
   if (isVirtualized) {
-    gridProps['aria-rowcount'] = state.collection.size;
+    gridProps['aria-rowcount'] = getItemCount(state.collection, (node) => node.type === 'item' || (node.type === 'loader' && node.props.isLoading));
     gridProps['aria-colcount'] = state.collection.columnCount;
   }
 

packages/@react-aria/gridlist/src/useGridList.ts

@@ -24,6 +24,7 @@ import {
   RefObject
 } from '@react-types/shared';
 import {filterDOMProps, mergeProps, useId} from '@react-aria/utils';
+import {getItemCount} from '@react-stately/collections';
 import {listMap} from './utils';
 import {ListState} from '@react-stately/list';
 import {useGridSelectionAnnouncement, useHighlightSelectionDescription} from '@react-aria/grid';
@@ -167,7 +168,7 @@ export function useGridList<T>(props: AriaGridListOptions<T>, state: ListState<T
   );
 
   if (isVirtualized) {
-    gridProps['aria-rowcount'] = state.collection.size;
+    gridProps['aria-rowcount'] = getItemCount(state.collection, (node) => node.type === 'item' || (node.type === 'loader' && node.props.isLoading));
     gridProps['aria-colcount'] = 1;
   }
 

packages/@react-aria/listbox/src/useOption.ts

@@ -121,7 +121,9 @@ export function useOption<T>(props: AriaOptionProps, state: ListState<T>, ref: R
   if (isVirtualized) {
     let index = Number(item?.index);
     optionProps['aria-posinset'] = Number.isNaN(index) ? undefined : index + 1;
-    optionProps['aria-setsize'] = getItemCount(state.collection);
+    // TODO: this is not great, but the loader sentinel is always in the collection even when loading isn't currently in progress.
+    // This same issue applies to the other collection elements, namely the row counts calculated for the top level parent element
+    optionProps['aria-setsize'] = getItemCount(state.collection, (node) => node.type === 'item' || (node.type === 'loader' && node.props.isLoading));
   }
 
   let onAction = data?.onAction ? () => data?.onAction?.(key) : undefined;

packages/@react-aria/menu/src/useMenuItem.ts

@@ -187,7 +187,7 @@ export function useMenuItem<T>(props: AriaMenuItemProps, state: TreeState<T>, re
 
   if (isVirtualized) {
     ariaProps['aria-posinset'] = item?.index;
-    ariaProps['aria-setsize'] = getItemCount(state.collection);
+    ariaProps['aria-setsize'] = getItemCount(state.collection, (node) => node.type === 'item' || (node.type === 'loader' && node.props.isLoading));
   }
 
   let onPressStart = (e: PressEvent) => {
@@ -307,8 +307,8 @@ export function useMenuItem<T>(props: AriaMenuItemProps, state: TreeState<T>, re
       ...mergeProps(
         domProps,
         linkProps,
-        isTrigger 
-          ? {onFocus: itemProps.onFocus, 'data-collection': itemProps['data-collection'], 'data-key': itemProps['data-key']} 
+        isTrigger
+          ? {onFocus: itemProps.onFocus, 'data-collection': itemProps['data-collection'], 'data-key': itemProps['data-key']}
           : itemProps,
         pressProps,
         hoverProps,

packages/@react-aria/selection/src/ListKeyboardDelegate.ts

@@ -71,14 +71,14 @@ export class ListKeyboardDelegate<T> implements KeyboardDelegate {
   }
 
   private isDisabled(item: Node<unknown>) {
-    return this.disabledBehavior === 'all' && (item.props?.isDisabled || this.disabledKeys.has(item.key));
+    return this.disabledBehavior === 'all' && (item.props?.isDisabled || this.disabledKeys.has(item.key) || (item.type === 'loader' && !item.props.isLoading));
   }
 
   private findNextNonDisabled(key: Key | null, getNext: (key: Key) => Key | null): Key | null {
     let nextKey = key;
     while (nextKey != null) {
       let item = this.collection.getItem(nextKey);
-      if (item?.type === 'loader' || (item?.type === 'item' && !this.isDisabled(item))) {
+      if ((item?.type === 'loader' || (item?.type === 'item')) && !this.isDisabled(item)) {
         return nextKey;
       }
 

packages/@react-aria/table/src/useTable.ts

@@ -86,7 +86,7 @@ export function useTable<T>(props: AriaTableProps, state: TableState<T> | TreeGr
 
   // Override to include header rows
   if (isVirtualized) {
-    gridProps['aria-rowcount'] = state.collection.size + state.collection.headerRows.length;
+    gridProps['aria-rowcount'] = gridProps['aria-rowcount']! + state.collection.headerRows.length;
   }
 
   if (tableNestedRows() && 'expandedKeys' in state) {

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

@@ -157,12 +157,13 @@ const iconStyles = style({
 });
 
 const loadingWrapperStyles = style({
+  ...focusRing(),
+  ...control({shape: 'default'}),
   gridColumnStart: '1',
   gridColumnEnd: '-1',
   display: 'flex',
   alignItems: 'center',
-  justifyContent: 'center',
-  marginY: 8
+  justifyContent: 'center'
 });
 
 const progressCircleStyles = style({

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

@@ -237,12 +237,13 @@ const iconStyles = style({
 });
 
 const loadingWrapperStyles = style({
+  ...focusRing(),
+  ...control({shape: 'default'}),
   gridColumnStart: '1',
   gridColumnEnd: '-1',
   display: 'flex',
   alignItems: 'center',
-  justifyContent: 'center',
-  marginY: 8
+  justifyContent: 'center'
 });
 
 const progressCircleStyles = style({

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

@@ -373,6 +373,22 @@ const centeredWrapper = style({
   height: 'full'
 });
 
+const loadingWrapperStyles = style({
+  width: 'full',
+  height: 'full',
+  boxSizing: 'border-box',
+  outlineStyle: {
+    default: 'none',
+    isFocusVisible: 'solid'
+  },
+  outlineWidth: 2,
+  outlineOffset: -2,
+  outlineColor: {
+    default: 'focus-ring',
+    forcedColors: 'Highlight'
+  }
+});
+
 export interface TableBodyProps<T> extends Omit<RACTableBodyProps<T>, 'style' | 'className'> {}
 
 /**
@@ -390,7 +406,10 @@ export const TableBody = /*#__PURE__*/ (forwardRef as forwardRefType)(function T
   // This is because we don't distinguish between loadingMore and loading in the layout, resulting in a different rect being used to build the body. Perhaps can be considered as a user error
   // if they pass loadingMore without having any other items in the table. Alternatively, could update the layout so it knows the current loading state.
   let loadMoreSpinner = (
-    <UNSTABLE_TableLoadingSentinel isLoading={loadingState === 'loadingMore'} onLoadMore={onLoadMore} className={style({height: 'full', width: 'full'})}>
+    <UNSTABLE_TableLoadingSentinel
+      isLoading={loadingState === 'loadingMore'}
+      onLoadMore={onLoadMore}
+      className={loadingWrapperStyles}>
       <div className={centeredWrapper}>
         <ProgressCircle
           isIndeterminate

packages/@react-stately/collections/src/getItemCount.ts

@@ -15,7 +15,7 @@ import {getChildNodes} from './getChildNodes';
 
 const cache = new WeakMap<Iterable<unknown>, number>();
 
-export function getItemCount<T>(collection: Collection<Node<T>>): number {
+export function getItemCount<T>(collection: Collection<Node<T>>, isValidItem: ((node: Node<T>) => boolean) = (node) => node.type === 'item'): number {
   let count = cache.get(collection);
   if (count != null) {
     return count;
@@ -25,11 +25,16 @@ export function getItemCount<T>(collection: Collection<Node<T>>): number {
   let counter = 0;
   let countItems = (items: Iterable<Node<T>>) => {
     for (let item of items) {
-      if (item.type === 'section') {
-        countItems(getChildNodes(item, collection));
-      } else if (item.type === 'item') {
+      if (isValidItem(item)) {
         counter++;
       }
+
+      // TODO: New collections vs old collection is different here. New collections for table will only return
+      // the header and body when the collection is iterated on, but old collections will return the body rows...
+      // I could forgo using getItemCount, but the collection.rows is different between old and new collection too... (old includes header row)
+      if (item.type === 'section' || item.type === 'tablebody') {
+        countItems(getChildNodes(item, collection));
+      }
     }
   };
 

packages/@react-stately/grid/src/useGridState.ts

@@ -59,55 +59,60 @@ export function useGridState<T extends object, C extends GridCollection<T>>(prop
   // Reset focused key if that item is deleted from the collection.
   const cachedCollection = useRef<C | null>(null);
   useEffect(() => {
-    if (selectionState.focusedKey != null && cachedCollection.current && !collection.getItem(selectionState.focusedKey)) {
-      const node = cachedCollection.current.getItem(selectionState.focusedKey);
-      const parentNode =
-        node?.parentKey != null && (node.type === 'cell' || node.type === 'rowheader' || node.type === 'column') ?
-        cachedCollection.current.getItem(node.parentKey) :
-        node;
-      if (!parentNode) {
-        selectionState.setFocusedKey(null);
-        return;
-      }
-      const cachedRows = cachedCollection.current.rows;
-      const rows = collection.rows;
-      const diff = cachedRows.length - rows.length;
-      let index = Math.min(
-        (
-          diff > 1 ?
-          Math.max(parentNode.index - diff + 1, 0) :
-          parentNode.index
-        ),
-        rows.length - 1);
-      let newRow: GridNode<T> | null = null;
-      while (index >= 0) {
-        if (!selectionManager.isDisabled(rows[index].key) && rows[index].type !== 'headerrow') {
-          newRow = rows[index];
-          break;
+    if (selectionState.focusedKey != null && cachedCollection.current) {
+      let item = collection.getItem(selectionState.focusedKey);
+      if (!item || (item.type === 'loader' && !item.props.isLoading)) {
+        const node = cachedCollection.current.getItem(selectionState.focusedKey);
+        const parentNode =
+          node?.parentKey != null && (node.type === 'cell' || node.type === 'rowheader' || node.type === 'column') ?
+          cachedCollection.current.getItem(node.parentKey) :
+          node;
+        if (!parentNode) {
+          selectionState.setFocusedKey(null);
+          return;
         }
-        // Find next, not disabled row.
-        if (index < rows.length - 1) {
-          index++;
-        // Otherwise, find previous, not disabled row.
-        } else {
-          if (index > parentNode.index) {
-            index = parentNode.index;
+        const cachedRows = cachedCollection.current.rows;
+        const rows = collection.rows;
+        const diff = cachedRows.length - rows.length;
+        let index = Math.min(
+          (
+            diff > 1 ?
+            Math.max(parentNode.index - diff + 1, 0) :
+            parentNode.index
+          ),
+          rows.length - 1);
+        let newRow: GridNode<T> | null = null;
+        while (index >= 0) {
+          // Find a row that is not disabled, nor is a header row, nor is a loader sentinel that isn't loading
+          if (!selectionManager.isDisabled(rows[index].key) && rows[index].type !== 'headerrow' && !(rows[index].type === 'loader' && !rows[index].props.isLoading)) {
+            newRow = rows[index];
+            break;
+          }
+          // Find next, not disabled row.
+          if (index < rows.length - 1) {
+            index++;
+          // Otherwise, find previous, not disabled row.
+          } else {
+            // eslint-disable-next-line max-depth
+            if (index > parentNode.index) {
+              index = parentNode.index;
+            }
+            index--;
           }
-          index--;
         }
-      }
-      if (newRow) {
-        const childNodes = newRow.hasChildNodes ? [...getChildNodes(newRow, collection)] : [];
-        const keyToFocus =
-          newRow.hasChildNodes &&
-          parentNode !== node &&
-          node &&
-          node.index < childNodes.length ?
-          childNodes[node.index].key :
-          newRow.key;
-        selectionState.setFocusedKey(keyToFocus);
-      } else {
-        selectionState.setFocusedKey(null);
+        if (newRow) {
+          const childNodes = newRow.hasChildNodes ? [...getChildNodes(newRow, collection)] : [];
+          const keyToFocus =
+            newRow.hasChildNodes &&
+            parentNode !== node &&
+            node &&
+            node.index < childNodes.length ?
+            childNodes[node.index].key :
+            newRow.key;
+          selectionState.setFocusedKey(keyToFocus);
+        } else {
+          selectionState.setFocusedKey(null);
+        }
       }
     }
     cachedCollection.current = collection;

packages/@react-stately/layout/src/ListLayout.ts

@@ -262,7 +262,7 @@ export class ListLayout<T, O extends ListLayoutOptions = ListLayoutOptions> exte
       y = 0;
     }
 
-    for (let node of collection) {
+    for (let node of collectionNodes) {
       let rowHeight = (this.rowHeight ?? this.estimatedRowHeight ?? DEFAULT_HEIGHT) + this.gap;
       // Skip rows before the valid rectangle unless they are already cached.
       if (node.type === 'item' && y + rowHeight < this.requestedRect.y && !this.isValid(node, y)) {
@@ -279,12 +279,13 @@ export class ListLayout<T, O extends ListLayoutOptions = ListLayoutOptions> exte
       }
 
       // Build each loader that exists in the collection that is outside the visible rect so that they are persisted
-      // at the proper estimated location
-      if (y > this.requestedRect.maxY) {
+      // at the proper estimated location. If the node.type is "section" then we don't do this shortcut since we have to
+      // build the sections to see how tall they are.
+      if ((node.type === 'item' || node.type === 'loader') && y > this.requestedRect.maxY) {
         let lastProcessedIndex = collectionNodes.indexOf(node);
         for (let loaderNode of loaderNodes) {
           let loaderNodeIndex = collectionNodes.indexOf(loaderNode);
-          // Subtract by an addition 1 since we've already added the current item's height to y
+          // Subtract by an additional 1 since we've already added the current item's height to y
           y += (loaderNodeIndex - lastProcessedIndex - 1) * rowHeight;
           let loader = this.buildChild(loaderNode, this.padding, y, null);
           nodes.push(loader);