Use flushSync to focus

src/DataGrid.tsx

@@ -318,7 +318,6 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
   const [isDragging, setDragging] = useState(false);
   const [draggedOverRowIdx, setOverRowIdx] = useState<number | undefined>(undefined);
   const [scrollToPosition, setScrollToPosition] = useState<PartialPosition | null>(null);
-  const [shouldFocusCell, setShouldFocusCell] = useState(false);
   const [previousRowIdx, setPreviousRowIdx] = useState(-1);
 
   const getColumnWidth = useCallback(
@@ -488,16 +487,6 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     latestDraggedOverRowIdx.current = rowIdx;
   }, []);
 
-  const focusCellOrCellContent = useCallback(() => {
-    const cell = getCellToScroll(gridRef.current!);
-    if (cell === null) return;
-
-    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 });
-  }, [gridRef]);
-
   /**
    * effects
    */
@@ -512,13 +501,6 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     }
   }, [selectedCellIsWithinSelectionBounds, selectedPosition]);
 
-  useLayoutEffect(() => {
-    if (shouldFocusCell) {
-      setShouldFocusCell(false);
-      focusCellOrCellContent();
-    }
-  }, [shouldFocusCell, focusCellOrCellContent]);
-
   useImperativeHandle(ref, () => ({
     element: gridRef.current,
     scrollToCell({ idx, rowIdx }) {
@@ -735,6 +717,16 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     );
   }
 
+  function focusCellOrCellContent() {
+    const cell = getCellToScroll(gridRef.current!);
+    if (cell === null) return;
+
+    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 });
+  }
+
   function selectCell(position: Position, enableEditor?: Maybe<boolean>): void {
     if (!isCellWithinSelectionBounds(position)) return;
     commitEditorChanges();
@@ -748,8 +740,10 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
       // Avoid re-renders if the selected cell state is the same
       scrollIntoView(getCellToScroll(gridRef.current!));
     } else {
-      setShouldFocusCell(true);
-      setSelectedPosition({ ...position, mode: 'SELECT' });
+      flushSync(() => {
+        setSelectedPosition({ ...position, mode: 'SELECT' });
+      });
+      focusCellOrCellContent();
     }
 
     if (onSelectedCellChange && !samePosition) {
@@ -897,6 +891,10 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     );
   }
 
+  function cancelEditing() {
+    setSelectedPosition(({ idx, rowIdx }) => ({ idx, rowIdx, mode: 'SELECT' }));
+  }
+
   function getCellEditor(rowIdx: number) {
     if (selectedPosition.rowIdx !== rowIdx || selectedPosition.mode === 'SELECT') return;
 
@@ -905,8 +903,12 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     const colSpan = getColSpan(column, lastFrozenColumnIndex, { type: 'ROW', row });
 
     const closeEditor = (shouldFocusCell: boolean) => {
-      setShouldFocusCell(shouldFocusCell);
-      setSelectedPosition(({ idx, rowIdx }) => ({ idx, rowIdx, mode: 'SELECT' }));
+      if (shouldFocusCell) {
+        flushSync(cancelEditing);
+        focusCellOrCellContent();
+      } else {
+        cancelEditing();
+      }
     };
 
     const onRowChange = (row: R, commitChanges: boolean, shouldFocusCell: boolean) => {
@@ -917,16 +919,19 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
         // SELECT and this results in onRowChange getting called twice.
         flushSync(() => {
           updateRow(column, selectedPosition.rowIdx, row);
-          closeEditor(shouldFocusCell);
+          cancelEditing();
         });
+        if (shouldFocusCell) {
+          focusCellOrCellContent();
+        }
       } else {
         setSelectedPosition((position) => ({ ...position, row }));
       }
     };
 
     if (rows[selectedPosition.rowIdx] !== selectedPosition.originalRow) {
       // Discard changes if rows are updated from outside
-      closeEditor(false);
+      cancelEditing();
     }
 
     return (

src/cellRenderers/renderCheckbox.tsx

@@ -24,6 +24,7 @@ const checkboxClassname = `rdg-checkbox-input ${checkbox}`;
 
 export function renderCheckbox({ onChange, indeterminate, ...props }: RenderCheckboxProps) {
   function handleChange(e: React.ChangeEvent<HTMLInputElement>) {
+    // https://github.com/facebook/react/issues/31358
     onChange(e.target.checked, (e.nativeEvent as MouseEvent).shiftKey);
   }