Fix cell stealing focus on click (#3790)

* Add a failing test

* Fix test

* Fix group selection

* Simplify

* Update test/browser/column/renderCell.test.tsx

Co-authored-by: Nicolas Stepien <[email protected]>

* Update test/browser/column/renderCell.test.tsx

---------

Co-authored-by: Nicolas Stepien <[email protected]>

Author: amanmahajan7

src/Cell.tsx

@@ -47,8 +47,8 @@ function Cell<R, SR>({
   );
   const isEditable = isCellEditableUtil(column, row);
 
-  function selectCellWrapper(openEditor?: boolean) {
-    selectCell({ rowIdx, idx: column.idx }, openEditor);
+  function selectCellWrapper(enableEditor?: boolean) {
+    selectCell({ rowIdx, idx: column.idx }, { enableEditor });
   }
 
   function handleMouseEvent(

src/DataGrid.tsx

@@ -57,6 +57,7 @@ import type {
   Position,
   Renderers,
   RowsChangeData,
+  SelectCellOptions,
   SelectHeaderRowEvent,
   SelectRowEvent,
   SortColumn
@@ -109,7 +110,7 @@ export type DefaultColumnOptions<R, SR> = Pick<
 export interface DataGridHandle {
   element: HTMLDivElement | null;
   scrollToCell: (position: PartialPosition) => void;
-  selectCell: (position: Position, enableEditor?: Maybe<boolean>) => void;
+  selectCell: (position: Position, options: SelectCellOptions) => void;
 }
 
 type SharedDivProps = Pick<
@@ -502,7 +503,7 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
   /**
    * callbacks
    */
-  const focusCellOrCellContent = useCallback(
+  const focusCell = useCallback(
     (shouldScroll = true) => {
       const cell = getCellToScroll(gridRef.current!);
       if (cell === null) return;
@@ -511,10 +512,7 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
         scrollIntoView(cell);
       }
 
-      // Focus cell content when available instead of the cell itself
-      const elementToFocus =
-        cell.querySelector<Element & HTMLOrSVGElement>('[tabindex="0"]') ?? cell;
-      elementToFocus.focus({ preventScroll: true });
+      cell.focus({ preventScroll: true });
     },
     [gridRef]
   );
@@ -528,11 +526,11 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
         focusSinkRef.current.focus({ preventScroll: true });
         scrollIntoView(focusSinkRef.current);
       } else {
-        focusCellOrCellContent();
+        focusCell();
       }
       setShouldFocusCell(false);
     }
-  }, [shouldFocusCell, focusCellOrCellContent, selectedPosition.idx]);
+  }, [shouldFocusCell, focusCell, selectedPosition.idx]);
 
   useImperativeHandle(ref, () => ({
     element: gridRef.current,
@@ -654,7 +652,7 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
   function handleFocus(event: React.FocusEvent<HTMLDivElement>) {
     // select the first header cell if the focus event is triggered by the grid
     if (event.target === event.currentTarget) {
-      selectHeaderCell({ idx: minColIdx, rowIdx: headerRowsCount });
+      selectHeaderCell({ idx: minColIdx, rowIdx: headerRowsCount }, { shouldFocusCell: true });
     }
   }
 
@@ -777,7 +775,7 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
 
   function handleDragHandleClick() {
     // keep the focus on the cell but do not scroll
-    focusCellOrCellContent(false);
+    focusCell(false);
   }
 
   function handleDragHandleDoubleClick(event: React.MouseEvent<HTMLDivElement>) {
@@ -838,20 +836,20 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     );
   }
 
-  function selectCell(position: Position, enableEditor?: Maybe<boolean>): void {
+  function selectCell(position: Position, options?: SelectCellOptions): void {
     if (!isCellWithinSelectionBounds(position)) return;
     commitEditorChanges();
 
     const samePosition = isSamePosition(selectedPosition, position);
 
-    if (enableEditor && isCellEditable(position)) {
+    if (options?.enableEditor && isCellEditable(position)) {
       const row = rows[position.rowIdx];
       setSelectedPosition({ ...position, mode: 'EDIT', row, originalRow: row });
     } else if (samePosition) {
       // Avoid re-renders if the selected cell state is the same
       scrollIntoView(getCellToScroll(gridRef.current!));
     } else {
-      setShouldFocusCell(true);
+      setShouldFocusCell(options?.shouldFocusCell === true);
       setSelectedPosition({ ...position, mode: 'SELECT' });
     }
 
@@ -864,8 +862,8 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     }
   }
 
-  function selectHeaderCell({ idx, rowIdx }: Position) {
-    selectCell({ rowIdx: minRowIdx + rowIdx - 1, idx });
+  function selectHeaderCell({ idx, rowIdx }: Position, options?: SelectCellOptions): void {
+    selectCell({ rowIdx: minRowIdx + rowIdx - 1, idx }, options);
   }
 
   function getNextPosition(key: string, ctrlKey: boolean, shiftKey: boolean): Position {
@@ -952,7 +950,7 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
       isCellWithinBounds: isCellWithinSelectionBounds
     });
 
-    selectCell(nextSelectedCellPosition);
+    selectCell(nextSelectedCellPosition, { shouldFocusCell: true });
   }
 
   function getDraggedOverCellIdx(currentRowIdx: number): number | undefined {

src/GroupRow.tsx

@@ -49,7 +49,7 @@ function GroupedRow<R, SR>({
   const idx = viewportColumns[0].key === SELECT_COLUMN_KEY ? row.level + 1 : row.level;
 
   function handleSelectGroup() {
-    selectCell({ rowIdx, idx: -1 });
+    selectCell({ rowIdx, idx: -1 }, { shouldFocusCell: true });
   }
 
   const selectionValue = useMemo(

src/hooks/useRovingTabIndex.ts

@@ -10,7 +10,13 @@ export function useRovingTabIndex(isSelected: boolean) {
   }
 
   function onFocus(event: React.FocusEvent<HTMLDivElement>) {
-    if (event.target !== event.currentTarget) {
+    const elementToFocus = event.currentTarget.querySelector<Element & HTMLOrSVGElement>(
+      '[tabindex="0"]'
+    );
+
+    // Focus cell content when available instead of the cell itself
+    if (elementToFocus !== null) {
+      elementToFocus.focus({ preventScroll: true });
       setIsChildFocused(true);
     }
   }

src/index.ts

@@ -48,6 +48,7 @@ export type {
   RenderSummaryCellProps,
   RowHeightArgs,
   RowsChangeData,
+  SelectCellOptions,
   SelectHeaderRowEvent,
   SelectRowEvent,
   SortColumn,

src/types.ts

@@ -167,7 +167,7 @@ interface BaseCellRendererProps<TRow, TSummaryRow = unknown>
       'onCellMouseDown' | 'onCellClick' | 'onCellDoubleClick' | 'onCellContextMenu'
     > {
   rowIdx: number;
-  selectCell: (position: Position, enableEditor?: Maybe<boolean>) => void;
+  selectCell: (position: Position, options?: SelectCellOptions) => void;
 }
 
 export interface CellRendererProps<TRow, TSummaryRow>
@@ -203,7 +203,7 @@ interface SelectCellKeyDownArgs<TRow, TSummaryRow = unknown> {
   column: CalculatedColumn<TRow, TSummaryRow>;
   row: TRow;
   rowIdx: number;
-  selectCell: (position: Position, enableEditor?: Maybe<boolean>) => void;
+  selectCell: (position: Position, options?: SelectCellOptions) => void;
 }
 
 export interface EditCellKeyDownArgs<TRow, TSummaryRow = unknown> {
@@ -334,6 +334,11 @@ export interface Renderers<TRow, TSummaryRow> {
   noRowsFallback?: Maybe<ReactNode>;
 }
 
+export interface SelectCellOptions {
+  enableEditor?: Maybe<boolean>;
+  shouldFocusCell?: Maybe<boolean>;
+}
+
 export interface ColumnWidth {
   readonly type: 'resized' | 'measured';
   readonly width: number;

test/browser/column/renderCell.test.tsx

@@ -116,6 +116,33 @@ describe('Custom cell renderer', () => {
   });
 });
 
+test('Focus child if it sets tabIndex', async () => {
+  const column: Column<Row> = {
+    key: 'test',
+    name: 'test',
+    renderCell(props) {
+      return (
+        <>
+          <button type="button" tabIndex={props.tabIndex}>
+            value: {props.row.id}
+          </button>
+          <span>External Text</span>
+        </>
+      );
+    }
+  };
+
+  page.render(<DataGrid columns={[column]} rows={[{ id: 1 }]} />);
+
+  const button = page.getByRole('button', { name: 'value: 1' });
+  await userEvent.click(page.getByText('External Text'));
+  expect(button).toHaveFocus();
+  await userEvent.tab();
+  expect(button).not.toHaveFocus();
+  await userEvent.click(button);
+  expect(button).toHaveFocus();
+});
+
 test('Cell should not steal focus when the focus is outside the grid and cell is recreated', async () => {
   const columns: readonly Column<Row>[] = [{ key: 'id', name: 'ID' }];