[DataGrid] Support state export and restore on grid density

packages/x-data-grid-pro/src/tests/statePersistence.DataGridPro.test.tsx

@@ -78,6 +78,7 @@ const FULL_INITIAL_STATE: GridInitialState = {
   sorting: {
     sortModel: [{ field: 'id', sort: 'desc' }],
   },
+  density: 'compact',
 };
 
 describe('<DataGridPro /> - State persistence', () => {
@@ -135,6 +136,7 @@ describe('<DataGridPro /> - State persistence', () => {
         sorting: {
           sortModel: [],
         },
+        density: 'standard',
       });
     });
 
@@ -165,6 +167,7 @@ describe('<DataGridPro /> - State persistence', () => {
           paginationMode="server"
           rowCount={FULL_INITIAL_STATE.pagination?.rowCount}
           pinnedColumns={FULL_INITIAL_STATE.pinnedColumns}
+          density={FULL_INITIAL_STATE.density}
           // Some portable states don't have a controllable model
           initialState={{
             columns: {
@@ -191,6 +194,7 @@ describe('<DataGridPro /> - State persistence', () => {
           paginationMode="server"
           rowCount={FULL_INITIAL_STATE.pagination?.rowCount}
           pinnedColumns={FULL_INITIAL_STATE.pinnedColumns}
+          density={FULL_INITIAL_STATE.density}
           // Some portable states don't have a controllable model
           initialState={{
             columns: {
@@ -226,6 +230,7 @@ describe('<DataGridPro /> - State persistence', () => {
         apiRef.current.setColumnIndex('category', 1);
         apiRef.current.setColumnWidth('category', 75);
         apiRef.current.setColumnVisibilityModel({ idBis: false });
+        apiRef.current.setDensity('compact');
       });
       expect(apiRef.current.exportState()).to.deep.equal(FULL_INITIAL_STATE);
     });

packages/x-data-grid/src/hooks/features/density/useGridDensity.tsx

@@ -7,6 +7,7 @@ import { GridDensityApi } from '../../../models/api/gridDensityApi';
 import { DataGridProcessedProps } from '../../../models/props/DataGridProps';
 import { gridDensitySelector } from './densitySelector';
 import { GridStateInitializer } from '../../utils/useGridInitializeState';
+import { GridPipeProcessor, useGridRegisterPipeProcessor } from '../../core/pipeProcessing';
 
 export const densityStateInitializer: GridStateInitializer<
   Pick<DataGridProcessedProps, 'initialState' | 'density'>
@@ -17,7 +18,7 @@ export const densityStateInitializer: GridStateInitializer<
 
 export const useGridDensity = (
   apiRef: React.MutableRefObject<GridPrivateApiCommunity>,
-  props: Pick<DataGridProcessedProps, 'density' | 'onDensityChange'>,
+  props: Pick<DataGridProcessedProps, 'density' | 'onDensityChange' | 'initialState'>,
 ): void => {
   const logger = useGridLogger(apiRef, 'useDensity');
 
@@ -29,6 +30,9 @@ export const useGridDensity = (
     changeEvent: 'densityChange',
   });
 
+  /**
+   * API METHODS
+   */
   const setDensity = useEventCallback<GridDensityApi['setDensity']>((newDensity): void => {
     const currentDensity = gridDensitySelector(apiRef.current.state);
     if (currentDensity === newDensity) {
@@ -43,15 +47,62 @@ export const useGridDensity = (
     }));
   });
 
-  React.useEffect(() => {
-    if (props.density) {
-      apiRef.current.setDensity(props.density);
-    }
-  }, [apiRef, props.density]);
-
   const densityApi: GridDensityApi = {
     setDensity,
   };
 
   useGridApiMethod(apiRef, densityApi, 'public');
+
+  /**
+   * PRE-PROCESSING
+   */
+  const stateExportPreProcessing = React.useCallback<GridPipeProcessor<'exportState'>>(
+    (prevState, context) => {
+      const exportedDensity = gridDensitySelector(apiRef.current.state);
+
+      const shouldExportRowCount =
+        // Always export if the `exportOnlyDirtyModels` property is not activated
+        !context.exportOnlyDirtyModels ||
+        // Always export if the `density` is controlled
+        props.density != null ||
+        // Always export if the `density` has been initialized
+        props.initialState?.density != null;
+
+      if (!shouldExportRowCount) {
+        return prevState;
+      }
+
+      return {
+        ...prevState,
+        density: exportedDensity,
+      };
+    },
+    [apiRef, props.density, props.initialState?.density],
+  );
+
+  const stateRestorePreProcessing = React.useCallback<GridPipeProcessor<'restoreState'>>(
+    (params, context) => {
+      const restoredDensity = context.stateToRestore?.density
+        ? context.stateToRestore.density
+        : gridDensitySelector(apiRef.current.state);
+      apiRef.current.setState((state) => ({
+        ...state,
+        density: restoredDensity,
+      }));
+      return params;
+    },
+    [apiRef],
+  );
+
+  useGridRegisterPipeProcessor(apiRef, 'exportState', stateExportPreProcessing);
+  useGridRegisterPipeProcessor(apiRef, 'restoreState', stateRestorePreProcessing);
+
+  /**
+   * EFFECTS
+   */
+  React.useEffect(() => {
+    if (props.density) {
+      apiRef.current.setDensity(props.density);
+    }
+  }, [apiRef, props.density]);
 };