[DataGridPro] Keep row children expansion when updating the rows

packages/grid/_modules_/grid/hooks/core/rowGroupsPerProcessing/gridRowGroupsPreProcessingApi.ts

@@ -3,6 +3,7 @@ import { GridRowTreeConfig, GridRowId, GridRowsLookup } from '../../../models/gr
 export interface GridRowGroupParams {
   ids: GridRowId[];
   idRowsLookup: GridRowsLookup;
+  previousTree: GridRowTreeConfig | null;
 }
 
 export interface GridRowGroupingResult {

packages/grid/_modules_/grid/hooks/features/rows/useGridRows.ts

@@ -10,6 +10,7 @@ import {
   GridRowsProp,
   GridRowIdGetter,
   GridRowTreeNodeConfig,
+  GridRowTreeConfig,
 } from '../../../models/gridRows';
 import { useGridApiMethod } from '../../utils/useGridApiMethod';
 import { useGridLogger } from '../../utils/useGridLogger';
@@ -24,8 +25,10 @@ import {
 import { GridSignature, useGridApiEventHandler } from '../../utils/useGridApiEventHandler';
 import { GridRowGroupParams } from '../../core/rowGroupsPerProcessing';
 
+type GridRowInternalCacheValue = Omit<GridRowGroupParams, 'previousTree'>;
+
 interface GridRowsInternalCacheState {
-  value: GridRowGroupParams;
+  value: GridRowInternalCacheValue;
 
   /**
    * The value of the properties used by the grouping when the internal cache was created
@@ -69,7 +72,7 @@ const convertGridRowsPropToState = ({
 }: ConvertGridRowsPropToStateParams): GridRowsInternalCacheState => {
   const props = inputProps ?? prevState.props;
 
-  let value: GridRowGroupParams;
+  let value: GridRowInternalCacheValue;
   if (rows) {
     value = {
       idRowsLookup: {},
@@ -94,14 +97,15 @@ const convertGridRowsPropToState = ({
 
 const getRowsStateFromCache = (
   rowsCache: GridRowsInternalCache,
+  previousTree: GridRowTreeConfig | null,
   apiRef: GridApiRef,
 ): GridRowsState => {
   const {
     props: { rowCount: propRowCount = 0 },
     value,
   } = rowsCache.state;
 
-  const groupingResponse = apiRef.current.unstable_groupRows(value);
+  const groupingResponse = apiRef.current.unstable_groupRows({ ...value, previousTree });
 
   const dataTopLevelRowCount = Object.values(groupingResponse.tree).filter(
     (node) => node.parent == null,
@@ -157,7 +161,7 @@ export const useGridRows = (
     });
     rowsCache.current.lastUpdateMs = Date.now();
 
-    return { ...state, rows: getRowsStateFromCache(rowsCache.current, apiRef) };
+    return { ...state, rows: getRowsStateFromCache(rowsCache.current, null, apiRef) };
   });
 
   const getRow = React.useCallback<GridRowApi['getRow']>(
@@ -172,7 +176,11 @@ export const useGridRows = (
         rowsCache.current.lastUpdateMs = Date.now();
         apiRef.current.setState((state) => ({
           ...state,
-          rows: getRowsStateFromCache(rowsCache.current, apiRef),
+          rows: getRowsStateFromCache(
+            rowsCache.current,
+            gridRowTreeSelector(apiRef.current.state),
+            apiRef,
+          ),
         }));
         apiRef.current.publishEvent(GridEvents.rowsSet);
         apiRef.current.forceUpdate();
@@ -244,7 +252,7 @@ export const useGridRows = (
 
       const deletedRowIds: GridRowId[] = [];
 
-      const newStateValue: GridRowGroupParams = {
+      const newStateValue: GridRowInternalCacheValue = {
         idRowsLookup: { ...rowsCache.current.state.value.idRowsLookup },
         ids: [...rowsCache.current.state.value.ids],
       };

packages/grid/_modules_/grid/utils/tree/buildRowTree.test.ts

@@ -18,6 +18,7 @@ describe('buildRowTree', () => {
         { id: 2, path: [{ key: 'C', field: null }] },
       ],
       defaultGroupingExpansionDepth: 0,
+      previousTree: null,
     });
 
     expect(
@@ -26,9 +27,9 @@ describe('buildRowTree', () => {
         childrenExpanded: node.childrenExpanded,
       })),
     ).to.deep.equal([
-      { id: 0, childrenExpanded: false },
-      { id: 1, childrenExpanded: false },
-      { id: 2, childrenExpanded: false },
+      { id: 0, childrenExpanded: undefined },
+      { id: 1, childrenExpanded: undefined },
+      { id: 2, childrenExpanded: undefined },
     ]);
   });
 
@@ -60,6 +61,7 @@ describe('buildRowTree', () => {
         },
       ],
       defaultGroupingExpansionDepth: 2,
+      previousTree: null,
     });
 
     expect(
@@ -70,7 +72,7 @@ describe('buildRowTree', () => {
     ).to.deep.equal([
       { id: 0, childrenExpanded: true },
       { id: 1, childrenExpanded: true },
-      { id: 2, childrenExpanded: false },
+      { id: 2, childrenExpanded: undefined },
     ]);
   });
 
@@ -102,6 +104,7 @@ describe('buildRowTree', () => {
         },
       ],
       defaultGroupingExpansionDepth: -1,
+      previousTree: null,
     });
 
     expect(
@@ -112,7 +115,7 @@ describe('buildRowTree', () => {
     ).to.deep.equal([
       { id: 0, childrenExpanded: true },
       { id: 1, childrenExpanded: true },
-      { id: 2, childrenExpanded: false },
+      { id: 2, childrenExpanded: undefined },
     ]);
   });
 
@@ -153,6 +156,7 @@ describe('buildRowTree', () => {
         },
       ],
       defaultGroupingExpansionDepth: 0,
+      previousTree: null,
     });
 
     expect(
@@ -206,6 +210,7 @@ describe('buildRowTree', () => {
         },
       ],
       defaultGroupingExpansionDepth: 0,
+      previousTree: null,
     });
 
     expect(response.treeDepth).to.equal(3);
@@ -231,6 +236,7 @@ describe('buildRowTree', () => {
         },
       ],
       defaultGroupingExpansionDepth: 0,
+      previousTree: null,
     });
 
     expect(response).to.deep.equal({
@@ -266,7 +272,7 @@ describe('buildRowTree', () => {
         1: {
           children: undefined,
           depth: 2,
-          childrenExpanded: false,
+          childrenExpanded: undefined,
           groupingField: null,
           groupingKey: 'A',
           id: 1,

packages/grid/_modules_/grid/utils/tree/buildRowTree.ts

@@ -71,8 +71,13 @@ export const buildRowTree = (params: BuildRowTreeParams): GridRowGroupingResult
   const groupingCriteriaToIdTree: GridGroupingCriteriaToIdTree = {};
 
   const isGroupExpandedByDefault = (node: GridRowTreeNodeConfig) => {
+    const previousExpansion = params.previousTree?.[node.id]?.childrenExpanded;
+    if (previousExpansion != null) {
+      return previousExpansion;
+    }
+
     if (!node.children || !node.children.length) {
-      return false;
+      return undefined;
     }
 
     if (params.isGroupExpandedByDefault) {

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

@@ -158,6 +158,17 @@ describe('<DataGridPro /> - Tree Data', () => {
       expect(getColumnHeadersTextContent()).to.deep.equal(['Group', 'nameBis']);
       expect(getColumnValues(1)).to.deep.equal(['1', '2']);
     });
+
+    it('should keep children expansion when changing some of the rows', () => {
+      const { setProps } = render(<Test rows={[{ name: 'A' }, { name: 'A.A' }]} />);
+      expect(getColumnValues(1)).to.deep.equal(['A']);
+      apiRef.current.setRowChildrenExpansion('A', true);
+      expect(getColumnValues(1)).to.deep.equal(['A', 'A.A']);
+      setProps({
+        rows: [{ name: 'A' }, { name: 'A.A' }, { name: 'B' }, { name: 'B.A' }],
+      });
+      expect(getColumnValues(1)).to.deep.equal(['A', 'A.A', 'B']);
+    });
   });
 
   describe('prop: getTreeDataPath', () => {