From 67c7be97e5754f6422bca66588c80656ee3db33b Mon Sep 17 00:00:00 2001
From: Kipre <Kipre>
Date: Wed, 22 Nov 2023 22:57:16 +0000
Subject: [PATCH] Recompute column widths when the number of columns changes

---
 src/hooks/useColumnWidths.ts         | 12 ++++--
 website/Nav.tsx                      |  3 ++
 website/demos/AddOrRemoveColumns.tsx | 63 ++++++++++++++++++++++++++++
 website/root.tsx                     |  2 +
 4 files changed, 77 insertions(+), 3 deletions(-)
 create mode 100644 website/demos/AddOrRemoveColumns.tsx

diff --git a/src/hooks/useColumnWidths.ts b/src/hooks/useColumnWidths.ts
index 5e6f7ed136..c5b34ba0fe 100644
--- a/src/hooks/useColumnWidths.ts
+++ b/src/hooks/useColumnWidths.ts
@@ -18,18 +18,24 @@ export function useColumnWidths<R, SR>(
   onColumnResize: DataGridProps<R, SR>['onColumnResize']
 ) {
   const prevGridWidthRef = useRef(gridWidth);
+  const prevNbOfColumns = useRef(columns.length);
+
   const columnsCanFlex: boolean = columns.length === viewportColumns.length;
+
+  const nbOfColumnsChanged = columns.length !== prevNbOfColumns.current;
   // Allow columns to flex again when...
   const ignorePreviouslyMeasuredColumns: boolean =
-    // there is enough space for columns to flex and the grid was resized
-    columnsCanFlex && gridWidth !== prevGridWidthRef.current;
+    // there is enough space for columns to flex and the grid was resized or
+    // the nb of columns changed
+    columnsCanFlex && (gridWidth !== prevGridWidthRef.current || nbOfColumnsChanged);
   const newTemplateColumns = [...templateColumns];
   const columnsToMeasure: string[] = [];
 
   for (const { key, idx, width } of viewportColumns) {
+    const previousMeasureIsRelevant = measuredColumnWidths.has(key) && !nbOfColumnsChanged;
     if (
       typeof width === 'string' &&
-      (ignorePreviouslyMeasuredColumns || !measuredColumnWidths.has(key)) &&
+      (ignorePreviouslyMeasuredColumns || !previousMeasureIsRelevant) &&
       !resizedColumnWidths.has(key)
     ) {
       newTemplateColumns[idx] = width;
diff --git a/website/Nav.tsx b/website/Nav.tsx
index 49fa4db6b6..5baa69beba 100644
--- a/website/Nav.tsx
+++ b/website/Nav.tsx
@@ -120,6 +120,9 @@ export default function Nav({ direction, onDirectionChange }: Props) {
       <NavLink to="/resizable-grid" end className={getActiveClassname}>
         Resizable Grid
       </NavLink>
+      <NavLink to="/add-remove-columns" end className={getActiveClassname}>
+        Add or Remove Columns
+      </NavLink>
       <NavLink to="/rows-reordering" end className={getActiveClassname}>
         Rows Reordering
       </NavLink>
diff --git a/website/demos/AddOrRemoveColumns.tsx b/website/demos/AddOrRemoveColumns.tsx
new file mode 100644
index 0000000000..14abba2f7b
--- /dev/null
+++ b/website/demos/AddOrRemoveColumns.tsx
@@ -0,0 +1,63 @@
+import { useMemo, useState } from 'react';
+
+import type { Column } from '../../src';
+import DataGrid from '../../src';
+import type { Props } from './types';
+
+interface Row {
+  label: string;
+  first: number;
+  second: number;
+  third: number;
+  fourth: number;
+  fifth: number;
+}
+
+const allColumns: Column<Row>[] = ['Label', 'First', 'Second', 'Third', 'Fourth', 'Fifth'].map(
+  (x) => ({ key: x.toLowerCase(), name: x })
+);
+
+const rows: Row[] = Array.from({ length: 8 }, (_, i) => ({
+  label: `Attribute #${i + 1}`,
+  first: 0.5 + Math.random() * 2,
+  second: 0.5 + Math.random() * 2,
+  third: 0.5 + Math.random() * 2,
+  fourth: 0.5 + Math.random() * 2,
+  fifth: 0.5 + Math.random() * 2
+}));
+
+export default function AddOrRemoveColumns({ direction }: Props) {
+  const [nbCols, setNbCols] = useState(4);
+  const columns = useMemo(() => allColumns.slice(0, nbCols), [nbCols]);
+
+  return (
+    <>
+      <div style={{ width: '300px', marginBlock: "1rem" }}>
+        <label htmlFor="nbcols">Choose the number of columns:</label>
+        <br />
+        <input
+          type="range"
+          id="nbcols"
+          list="markers"
+          min="1"
+          max="5"
+          step="2"
+          onChange={(e) => setNbCols(parseInt(e.target.value, 10))}
+        />
+
+        <datalist id="markers">
+          <option value="1" />
+          <option value="3" />
+          <option value="5" />
+        </datalist>
+      </div>
+
+      <DataGrid
+        columns={columns}
+        rows={rows}
+        style={{ resize: 'both', width: '800px', height: '3OOpx' }}
+        direction={direction}
+      />
+    </>
+  );
+}
diff --git a/website/root.tsx b/website/root.tsx
index 0201879cd1..cc532db401 100644
--- a/website/root.tsx
+++ b/website/root.tsx
@@ -27,6 +27,7 @@ import ScrollToCell from './demos/ScrollToCell';
 import TreeView from './demos/TreeView';
 import VariableRowHeight from './demos/VariableRowHeight';
 import Nav from './Nav';
+import AddOrRemoveColumns from './demos/AddOrRemoveColumns';
 
 const mainClassname = css`
   display: flex;
@@ -63,6 +64,7 @@ function Root() {
           <Route path="million-cells" element={<MillionCells direction={direction} />} />
           <Route path="no-rows" element={<NoRows direction={direction} />} />
           <Route path="resizable-grid" element={<ResizableGrid direction={direction} />} />
+          <Route path="add-remove-columns" element={<AddOrRemoveColumns direction={direction} />} />
           <Route path="rows-reordering" element={<RowsReordering direction={direction} />} />
           <Route path="scroll-to-cell" element={<ScrollToCell direction={direction} />} />
           <Route path="tree-view" element={<TreeView direction={direction} />} />