Add columnWidths and onColumnWidthsChange props

Author: amanmahajan7

src/DataGrid.tsx

@@ -109,7 +109,6 @@ export interface DataGridHandle {
   element: HTMLDivElement | null;
   scrollToCell: (position: PartialPosition) => void;
   selectCell: (position: Position, enableEditor?: Maybe<boolean>) => void;
-  resetMeasuredColumnWidths: () => void;
 }
 
 type SharedDivProps = Pick<
@@ -236,6 +235,9 @@ export interface DataGridProps<R, SR = unknown, K extends Key = Key> extends Sha
   direction?: Maybe<Direction>;
   'data-testid'?: Maybe<string>;
   'data-cy'?: Maybe<string>;
+
+  columnWidths?: Maybe<ReadonlyMap<string, number>>;
+  onColumnWidthsChange?: Maybe<(widths: ReadonlyMap<string, number>) => void>;
 }
 
 /**
@@ -295,7 +297,10 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
     'aria-describedby': ariaDescribedBy,
     'aria-rowcount': rawAriaRowCount,
     'data-testid': testId,
-    'data-cy': dataCy
+    'data-cy': dataCy,
+
+    columnWidths,
+    onColumnWidthsChange
   } = props;
 
   /**
@@ -324,9 +329,17 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
   const [resizedColumnWidths, setResizedColumnWidths] = useState(
     (): ReadonlyMap<string, number> => new Map()
   );
-  const [measuredColumnWidths, setMeasuredColumnWidths] = useState(
-    (): ReadonlyMap<string, number> => new Map()
+  const [measuredColumnWidthsInternal, setMeasuredColumnWidthsInternal] = useState(
+    (): ReadonlyMap<string, number> => columnWidths ?? new Map()
   );
+  const isColumnWidthsControlled = columnWidths != null && onColumnWidthsChange != null;
+  const measuredColumnWidths = isColumnWidthsControlled
+    ? columnWidths
+    : measuredColumnWidthsInternal;
+  const setMeasuredColumnWidths = isColumnWidthsControlled
+    ? onColumnWidthsChange
+    : setMeasuredColumnWidthsInternal;
+
   const [isDragging, setDragging] = useState(false);
   const [draggedOverRowIdx, setOverRowIdx] = useState<number | undefined>(undefined);
   const [scrollToPosition, setScrollToPosition] = useState<PartialPosition | null>(null);
@@ -541,10 +554,7 @@ export function DataGrid<R, SR = unknown, K extends Key = Key>(props: DataGridPr
         setScrollToPosition({ idx: scrollToIdx, rowIdx: scrollToRowIdx });
       }
     },
-    selectCell,
-    resetMeasuredColumnWidths() {
-      setMeasuredColumnWidths(new Map());
-    }
+    selectCell
   }));
 
   /**

src/hooks/useColumnWidths.ts

@@ -13,14 +13,18 @@ export function useColumnWidths<R, SR>(
   resizedColumnWidths: ReadonlyMap<string, number>,
   measuredColumnWidths: ReadonlyMap<string, number>,
   setResizedColumnWidths: StateSetter<ReadonlyMap<string, number>>,
-  setMeasuredColumnWidths: StateSetter<ReadonlyMap<string, number>>,
+  setMeasuredColumnWidths: (columnWidths: ReadonlyMap<string, number>) => void,
   onColumnResize: DataGridProps<R, SR>['onColumnResize']
 ) {
   const [columnToAutoResize, setColumnToAutoResize] = useState<{
     readonly key: string;
     readonly width: ResizedWidth;
   } | null>(null);
   const [prevGridWidth, setPreviousGridWidth] = useState(gridWidth);
+  const [temporaryMeasuredColumnWidths, setTemporaryMeasuredColumnWidths] = useState<ReadonlyMap<
+    string,
+    number
+  > | null>(null);
   const columnsCanFlex: boolean = columns.length === viewportColumns.length;
   // Allow columns to flex again when...
   const ignorePreviouslyMeasuredColumns: boolean =
@@ -38,7 +42,10 @@ export function useColumnWidths<R, SR>(
       columnsToMeasure.push(key);
     } else if (
       typeof width === 'string' &&
-      (ignorePreviouslyMeasuredColumns || !measuredColumnWidths.has(key)) &&
+      (ignorePreviouslyMeasuredColumns ||
+        (temporaryMeasuredColumnWidths !== null
+          ? temporaryMeasuredColumnWidths.has(key)
+          : !measuredColumnWidths.has(key))) &&
       !resizedColumnWidths.has(key)
     ) {
       newTemplateColumns[idx] = width;
@@ -54,22 +61,22 @@ export function useColumnWidths<R, SR>(
     setPreviousGridWidth(gridWidth);
     if (columnsToMeasure.length === 0) return;
 
-    setMeasuredColumnWidths((measuredColumnWidths) => {
-      const newMeasuredColumnWidths = new Map(measuredColumnWidths);
-      let hasChanges = false;
-
-      for (const key of columnsToMeasure) {
-        const measuredWidth = measureColumnWidth(gridRef, key);
-        hasChanges ||= measuredWidth !== measuredColumnWidths.get(key);
-        if (measuredWidth === undefined) {
-          newMeasuredColumnWidths.delete(key);
-        } else {
-          newMeasuredColumnWidths.set(key, measuredWidth);
-        }
+    const newMeasuredColumnWidths = new Map(measuredColumnWidths);
+    let hasChanges = false;
+
+    for (const key of columnsToMeasure) {
+      const measuredWidth = measureColumnWidth(gridRef, key);
+      hasChanges ||= measuredWidth !== measuredColumnWidths.get(key);
+      if (measuredWidth === undefined) {
+        newMeasuredColumnWidths.delete(key);
+      } else {
+        newMeasuredColumnWidths.set(key, measuredWidth);
       }
+    }
 
-      return hasChanges ? newMeasuredColumnWidths : measuredColumnWidths;
-    });
+    if (hasChanges) {
+      setMeasuredColumnWidths(newMeasuredColumnWidths);
+    }
 
     if (columnToAutoResize !== null) {
       const resizingKey = columnToAutoResize.key;
@@ -93,15 +100,14 @@ export function useColumnWidths<R, SR>(
     flushSync(() => {
       if (columnsCanFlex) {
         // remeasure all the columns that can flex and are not resized by the user
-        setMeasuredColumnWidths((measuredColumnWidths) => {
-          const newMeasuredColumnWidths = new Map(measuredColumnWidths);
-          for (const { key, width } of viewportColumns) {
-            if (resizingKey !== key && typeof width === 'string' && !resizedColumnWidths.has(key)) {
-              newMeasuredColumnWidths.delete(key);
-            }
+        const newMeasuredColumnWidths = new Map(measuredColumnWidths);
+        for (const { key, width } of viewportColumns) {
+          if (resizingKey !== key && typeof width === 'string' && !resizedColumnWidths.has(key)) {
+            newMeasuredColumnWidths.delete(key);
           }
-          return newMeasuredColumnWidths;
-        });
+        }
+
+        setTemporaryMeasuredColumnWidths(newMeasuredColumnWidths);
       }
 
       setColumnToAutoResize({
@@ -110,6 +116,8 @@ export function useColumnWidths<R, SR>(
       });
     });
 
+    setTemporaryMeasuredColumnWidths(null);
+
     if (onColumnResize) {
       const previousWidth = resizedColumnWidths.get(resizingKey);
       const newWidth =

website/routes/CommonFeatures.tsx

@@ -323,6 +323,9 @@ function CommonFeatures() {
   const [isExporting, setIsExporting] = useState(false);
   const gridRef = useRef<DataGridHandle>(null);
   const columns = useMemo(() => getColumns(countries, direction), [direction]);
+  const [columnWidths, setColumnWidths] = useState(
+    (): ReadonlyMap<string, number> => new Map<string, number>()
+  );
 
   const summaryRows = useMemo((): readonly SummaryRow[] => {
     return [
@@ -382,6 +385,9 @@ function CommonFeatures() {
         <button type="button" onClick={handleExportToPdf}>
           Export to PDF
         </button>
+        <button type="button" onClick={() => setColumnWidths(new Map())}>
+          Reset Width
+        </button>
       </div>
       <DataGrid
         ref={gridRef}
@@ -402,6 +408,8 @@ function CommonFeatures() {
         className="fill-grid"
         direction={direction}
         enableVirtualization={!isExporting}
+        columnWidths={columnWidths}
+        onColumnWidthsChange={setColumnWidths}
       />
     </>
   );