bug(data frame): Use sorted row data to find next row when hitting (shift+)enter

js/dataframe/cell.tsx

@@ -1,4 +1,4 @@
-import { flexRender } from "@tanstack/react-table";
+import { RowModel, flexRender } from "@tanstack/react-table";
 import { VirtualItem } from "@tanstack/react-virtual";
 import { Cell } from "@tanstack/table-core";
 import React, {
@@ -43,35 +43,34 @@ const CellStateClassEnum = {
 export type CellState = keyof typeof CellStateEnum;
 
 interface TableBodyCellProps {
-  id: string | null;
+  key: string;
+  rowId: string;
   containerRef: React.RefObject<HTMLDivElement>;
   cell: Cell<unknown[], unknown>;
   patchInfo: PatchInfo;
   columns: readonly string[];
   rowIndex: number;
   columnIndex: number;
   editCellsIsAllowed: boolean;
-  virtualRows: VirtualItem[];
+  getSortedRowModel: () => RowModel<unknown[]>;
   setData: (fn: (draft: unknown[][]) => void) => void;
   cellEditInfo: CellEdit | undefined;
   setCellEditMapAtLoc: SetCellEditMapAtLoc;
-  maxRowSize: number;
 }
 
 export const TableBodyCell: FC<TableBodyCellProps> = ({
-  id,
   containerRef,
+  rowId,
   cell,
   patchInfo,
   columns,
   rowIndex,
   columnIndex,
   editCellsIsAllowed,
-  virtualRows,
+  getSortedRowModel,
   cellEditInfo,
   setData,
   setCellEditMapAtLoc,
-  maxRowSize,
 }) => {
   const initialValue = cell.getValue();
 
@@ -147,13 +146,15 @@ export const TableBodyCell: FC<TableBodyCellProps> = ({
     const hasShift = e.shiftKey;
 
     const newColumnIndex = columnIndex! + (hasShift ? -1 : 1);
+
+    // Submit changes to the current cell
+    attemptUpdate();
+
     if (newColumnIndex < 0 || newColumnIndex >= columns.length) {
       // If the new column index is out of bounds, quit
       return;
     }
 
-    attemptUpdate();
-
     // Turn on editing in next cell!
     setCellEditMapAtLoc(rowIndex, newColumnIndex, (obj_draft) => {
       obj_draft.isEditing = true;
@@ -167,16 +168,26 @@ export const TableBodyCell: FC<TableBodyCellProps> = ({
 
     const hasShift = e.shiftKey;
 
-    const newRowIndex = rowIndex! + (hasShift ? -1 : 1);
-    if (newRowIndex < 0 || newRowIndex >= maxRowSize) {
-      // If the new row index is out of bounds, quit
+    const rowModel = getSortedRowModel();
+    const sortedRowIndex = rowModel.rows.findIndex((row) => row.id === rowId);
+    // Couldn't find row... silently quit
+    if (sortedRowIndex < 0) {
       return;
     }
+    const nextSortedRowIndex = sortedRowIndex! + (hasShift ? -1 : 1);
 
+    // Submit changes to the current cell
     attemptUpdate();
 
+    if (nextSortedRowIndex < 0 || nextSortedRowIndex >= rowModel.rows.length) {
+      // If the new row index is out of bounds, quit
+      return;
+    }
+
     // Turn on editing in the next cell!
-    setCellEditMapAtLoc(newRowIndex, columnIndex, (obj_draft) => {
+    // Get the original row index
+    const targetRowIndex = rowModel.rows[nextSortedRowIndex].index;
+    setCellEditMapAtLoc(targetRowIndex, columnIndex, (obj_draft) => {
       obj_draft.isEditing = true;
     });
   };
@@ -213,7 +224,6 @@ export const TableBodyCell: FC<TableBodyCellProps> = ({
     // Update the data!
     // updateCellsData updates the underlying data via `setData` and `setCellEditMapAtLoc`
     updateCellsData({
-      id,
       patchInfo: patchInfo,
       patches: [{ rowIndex, columnIndex, value: editValue }],
       onSuccess: (_patches) => {
@@ -238,7 +248,6 @@ export const TableBodyCell: FC<TableBodyCellProps> = ({
     initialValue,
     editValue,
     resetEditing,
-    id,
     patchInfo,
     columns,
     setData,
@@ -370,7 +379,7 @@ export const TableBodyCell: FC<TableBodyCellProps> = ({
 
   return (
     <td
-      key={cell.id}
+      id={cell.id}
       onClick={onClick}
       title={cellTitle}
       className={tableCellClass}

js/dataframe/data-update.tsx

@@ -22,7 +22,6 @@ export type CellPatchPy = {
 };
 
 export function updateCellsData({
-  id,
   patchInfo,
   patches,
   onSuccess,
@@ -31,7 +30,6 @@ export function updateCellsData({
   setData,
   setCellEditMapAtLoc,
 }: {
-  id: string | null;
   patchInfo: PatchInfo;
   patches: CellPatch[];
   onSuccess: (values: CellPatch[]) => void;

js/dataframe/index.tsx

@@ -392,8 +392,6 @@ const ShinyDataGrid: FC<ShinyDataGridProps<unknown>> = ({
     className += " html-fill-item";
   }
 
-  const maxRowSize = table.getSortedRowModel().rows.length;
-
   return (
     <>
       <div
@@ -490,28 +488,19 @@ const ShinyDataGrid: FC<ShinyDataGridProps<unknown>> = ({
 
                       return (
                         <TableBodyCell
-                          id={id}
-                          containerRef={containerRef}
                           key={cell.id}
+                          rowId={cell.row.id}
+                          containerRef={containerRef}
                           cell={cell}
                           patchInfo={patchInfo}
                           editCellsIsAllowed={editCellsIsAllowed}
                           columns={columns}
                           rowIndex={rowIndex}
                           columnIndex={columnIndex}
-                          // isEditing={
-                          //   editRowIndex === rowIndex &&
-                          //   editColumnIndex === columnIndex
-                          // }
-                          // setEditRowIndex={setEditRowIndex}
-                          // setEditColumnIndex={setEditColumnIndex}
-                          virtualRows={virtualRows}
-                          // cellEditMap={cellEditMap}
+                          getSortedRowModel={table.getSortedRowModel}
                           cellEditInfo={cellEditInfo}
-                          maxRowSize={maxRowSize}
                           setData={setData}
                           setCellEditMapAtLoc={setCellEditMapAtLoc}
-                          // updateCellsData={updateCellsData}
                         ></TableBodyCell>
                       );
                     })}