From 5dc5bd4b04eecf21c45cf6ea1169727ba5ab3be3 Mon Sep 17 00:00:00 2001
From: Flavien DELANGLE <flaviendelangle@gmail.com>
Date: Mon, 11 Apr 2022 10:09:08 +0200
Subject: [PATCH] [data grid] Fix column dimensions import/export with flex and
 resizing (#4311)

---
 .../columnResize/useGridColumnResize.tsx      |  2 +-
 .../statePersistence.DataGridPro.test.tsx     |  2 +-
 .../features/columns/gridColumnsInterfaces.ts |  2 +-
 .../features/columns/gridColumnsUtils.ts      | 27 +++++++++++++------
 .../hooks/features/columns/useGridColumns.tsx |  6 ++++-
 5 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/packages/grid/x-data-grid-pro/src/hooks/features/columnResize/useGridColumnResize.tsx b/packages/grid/x-data-grid-pro/src/hooks/features/columnResize/useGridColumnResize.tsx
index bf86d9d1b9a7..6a15c389e4a6 100644
--- a/packages/grid/x-data-grid-pro/src/hooks/features/columnResize/useGridColumnResize.tsx
+++ b/packages/grid/x-data-grid-pro/src/hooks/features/columnResize/useGridColumnResize.tsx
@@ -136,7 +136,7 @@ export const useGridColumnResize = (
 
     colDefRef.current!.computedWidth = newWidth;
     colDefRef.current!.width = newWidth;
-    colDefRef.current!.flex = undefined;
+    colDefRef.current!.flex = 0;
 
     colElementRef.current!.style.width = `${newWidth}px`;
     colElementRef.current!.style.minWidth = `${newWidth}px`;
diff --git a/packages/grid/x-data-grid-pro/src/tests/statePersistence.DataGridPro.test.tsx b/packages/grid/x-data-grid-pro/src/tests/statePersistence.DataGridPro.test.tsx
index 888c3041bce1..87c111db3cad 100644
--- a/packages/grid/x-data-grid-pro/src/tests/statePersistence.DataGridPro.test.tsx
+++ b/packages/grid/x-data-grid-pro/src/tests/statePersistence.DataGridPro.test.tsx
@@ -51,7 +51,7 @@ const FULL_INITIAL_STATE: GridInitialState = {
     dimensions: {
       idBis: {
         width: 75,
-        maxWidth: Infinity,
+        maxWidth: -1,
         minWidth: 50,
         flex: undefined,
       },
diff --git a/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsInterfaces.ts b/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsInterfaces.ts
index 401b312a8a56..619bb3124a59 100644
--- a/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsInterfaces.ts
+++ b/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsInterfaces.ts
@@ -19,7 +19,7 @@ export interface GridColumnsState {
   columnVisibilityModel: GridColumnVisibilityModel;
 }
 
-export type GridColumnDimensions = Pick<GridStateColDef, GridColumnDimensionProperties>;
+export type GridColumnDimensions = { [key in GridColumnDimensionProperties]?: number };
 
 export interface GridColumnsInitialState {
   columnVisibilityModel?: GridColumnVisibilityModel;
diff --git a/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts b/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts
index 0ba15406ed1b..1b53b6271743 100644
--- a/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts
+++ b/packages/grid/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts
@@ -53,7 +53,7 @@ export function computeFlexColumnsWidth({
   totalFlexUnits: number;
   flexColumns: {
     field: GridColDef['field'];
-    flex?: number;
+    flex?: number | null;
     minWidth?: number;
     maxWidth?: number;
   }[];
@@ -266,11 +266,16 @@ export const applyInitialState = (
   for (let i = 0; i < columnsWithUpdatedDimensions.length; i += 1) {
     const field = columnsWithUpdatedDimensions[i];
 
-    newColumnLookup[field] = {
+    const newColDef: Omit<GridStateColDef, 'computedWidth'> = {
       ...newColumnLookup[field],
-      ...dimensions[field],
       hasBeenResized: true,
     };
+
+    Object.entries(dimensions[field]).forEach(([key, value]) => {
+      newColDef[key as GridColumnDimensionProperties] = value === -1 ? Infinity : value;
+    });
+
+    newColumnLookup[field] = newColDef;
   }
 
   const newColumnsState: Omit<GridColumnsRawState, 'columnVisibilityModel'> = {
@@ -376,16 +381,22 @@ export const createColumnsState = ({
       columnsStateWithoutColumnVisibilityModel.all.push(field);
     }
 
-    let hasValidDimension = false;
-    if (!existingState.hasBeenResized) {
-      hasValidDimension = COLUMNS_DIMENSION_PROPERTIES.some((key) => newColumn[key] !== undefined);
-    }
+    let hasBeenResized = existingState.hasBeenResized;
+    COLUMNS_DIMENSION_PROPERTIES.forEach((key) => {
+      if (newColumn[key] !== undefined) {
+        hasBeenResized = true;
+
+        if (newColumn[key] === -1) {
+          newColumn[key] = Infinity;
+        }
+      }
+    });
 
     columnsStateWithoutColumnVisibilityModel.lookup[field] = {
       ...existingState,
       hide: newColumn.hide == null ? false : newColumn.hide,
       ...newColumn,
-      hasBeenResized: existingState.hasBeenResized || hasValidDimension,
+      hasBeenResized,
     };
   });
 
diff --git a/packages/grid/x-data-grid/src/hooks/features/columns/useGridColumns.tsx b/packages/grid/x-data-grid/src/hooks/features/columns/useGridColumns.tsx
index a0ae8d255491..92418a0e1bd8 100644
--- a/packages/grid/x-data-grid/src/hooks/features/columns/useGridColumns.tsx
+++ b/packages/grid/x-data-grid/src/hooks/features/columns/useGridColumns.tsx
@@ -323,7 +323,11 @@ export function useGridColumns(
         if (colDef.hasBeenResized) {
           const colDefDimensions: GridColumnDimensions = {};
           COLUMNS_DIMENSION_PROPERTIES.forEach((propertyName) => {
-            colDefDimensions[propertyName] = colDef[propertyName];
+            let propertyValue: number | undefined = colDef[propertyName];
+            if (propertyValue === Infinity) {
+              propertyValue = -1;
+            }
+            colDefDimensions[propertyName] = propertyValue;
           });
           dimensions[colDef.field] = colDefDimensions;
         }