[DataGridPro] Retain expansion state on rows update

docs/data/data-grid/tree-data/TreeDataFullExample.js

@@ -1,17 +1,23 @@
 import * as React from 'react';
-import { DataGridPro } from '@mui/x-data-grid-pro';
+import { DataGridPro, useGridApiRef } from '@mui/x-data-grid-pro';
 import { useDemoData } from '@mui/x-data-grid-generator';
+import Button from '@mui/material/Button';
 
 export default function TreeDataFullExample() {
+  const apiRef = useGridApiRef();
   const { data, loading } = useDemoData({
     dataSet: 'Employee',
     rowLength: 1000,
     treeData: { maxDepth: 2, groupingField: 'name', averageChildren: 200 },
   });
 
   return (
-    <div style={{ height: 400, width: '100%' }}>
-      <DataGridPro loading={loading} {...data} />
+    <div style={{ width: '100%' }}>
+      <Button onClick={() => apiRef.current?.expandAllRows()}>Expand all</Button>
+      <Button onClick={() => apiRef.current?.collapseAllRows()}>Collapse all</Button>
+      <div style={{ height: 400 }}>
+        <DataGridPro loading={loading} {...data} apiRef={apiRef} />
+      </div>
     </div>
   );
 }

docs/data/data-grid/tree-data/TreeDataFullExample.tsx

@@ -1,17 +1,23 @@
 import * as React from 'react';
-import { DataGridPro } from '@mui/x-data-grid-pro';
+import { DataGridPro, useGridApiRef } from '@mui/x-data-grid-pro';
 import { useDemoData } from '@mui/x-data-grid-generator';
+import Button from '@mui/material/Button';
 
 export default function TreeDataFullExample() {
+  const apiRef = useGridApiRef();
   const { data, loading } = useDemoData({
     dataSet: 'Employee',
     rowLength: 1000,
     treeData: { maxDepth: 2, groupingField: 'name', averageChildren: 200 },
   });
 
   return (
-    <div style={{ height: 400, width: '100%' }}>
-      <DataGridPro loading={loading} {...data} />
+    <div style={{ width: '100%' }}>
+      <Button onClick={() => apiRef.current?.expandAllRows()}>Expand all</Button>
+      <Button onClick={() => apiRef.current?.collapseAllRows()}>Collapse all</Button>
+      <div style={{ height: 400 }}>
+        <DataGridPro loading={loading} {...data} apiRef={apiRef} />
+      </div>
     </div>
   );
 }

docs/data/data-grid/tree-data/TreeDataFullExample.tsx.preview

@@ -1 +1,5 @@
-<DataGridPro loading={loading} {...data} />
+<Button onClick={() => apiRef.current?.expandAllRows()}>Expand all</Button>
+<Button onClick={() => apiRef.current?.collapseAllRows()}>Collapse all</Button>
+<div style={{ height: 400 }}>
+  <DataGridPro loading={loading} {...data} apiRef={apiRef} />
+</div>

docs/pages/x/api/data-grid/grid-api.json

@@ -26,11 +26,21 @@
       "type": { "description": "(options?: GridAutosizeOptions) => Promise<void>" },
       "required": true
     },
+    "collapseAllRows": {
+      "type": { "description": "() => void" },
+      "required": true,
+      "isProPlan": true
+    },
     "dataSource": { "type": { "description": "GridDataSourceApiBase" }, "required": true },
     "deleteFilterItem": {
       "type": { "description": "(item: GridFilterItem) => void" },
       "required": true
     },
+    "expandAllRows": {
+      "type": { "description": "() => void" },
+      "required": true,
+      "isProPlan": true
+    },
     "exportDataAsCsv": {
       "type": { "description": "(options?: GridCsvExportOptions) => void" },
       "required": true

docs/translations/api-docs/data-grid/grid-api.json

@@ -7,10 +7,12 @@
     "autosizeColumns": {
       "description": "Auto-size the columns of the grid based on the cells' content and the space available."
     },
+    "collapseAllRows": { "description": "Collapse all rows. Works for the client side data only." },
     "dataSource": { "description": "The data source API." },
     "deleteFilterItem": {
       "description": "Deletes a <a href=\"/x/api/data-grid/grid-filter-item/\">GridFilterItem</a>."
     },
+    "expandAllRows": { "description": "Expand all rows. Works for the client side data only." },
     "exportDataAsCsv": { "description": "Downloads and exports a CSV of the grid&#39;s data." },
     "exportDataAsExcel": {
       "description": "Downloads and exports an Excel file of the grid&#39;s data."
@@ -205,7 +207,9 @@
     "setQuickFilterValues": {
       "description": "Set the quick filter values to the one given by <code>values</code>"
     },
-    "setRowChildrenExpansion": { "description": "Expand or collapse a row children." },
+    "setRowChildrenExpansion": {
+      "description": "Expand or collapse a row children.<br />Only works for the client side data or to collapse already fetched server side rows.<br />For server-side data, use <code>dataSource.fetchRows(childId)</code> to fetch and expand the children."
+    },
     "setRowCount": { "description": "Sets the <code>rowCount</code> to a new value." },
     "setRowGroupingCriteriaIndex": {
       "description": "Sets the grouping index of a grouping criteria."

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

@@ -309,6 +309,45 @@ describe('<DataGridPro /> - Tree data', () => {
     });
   });
 
+  describe('apiRef: expandAllRows', () => {
+    it('should expand all rows', async () => {
+      render(<Test />);
+      expect(getColumnValues(1)).to.deep.equal(['A', 'B', 'C']);
+      act(() => apiRef.current?.expandAllRows());
+
+      expect(getColumnValues(1)).to.deep.equal([
+        'A',
+        'A.A',
+        'A.B',
+        'B',
+        'B.A',
+        'B.B',
+        'B.B.A',
+        'B.B.A.A',
+        'C',
+      ]);
+    });
+  });
+
+  describe('apiRef: collapseAllRows', () => {
+    it('should collapse all rows', () => {
+      render(<Test defaultGroupingExpansionDepth={-1} />);
+      expect(getColumnValues(1)).to.deep.equal([
+        'A',
+        'A.A',
+        'A.B',
+        'B',
+        'B.A',
+        'B.B',
+        'B.B.A',
+        'B.B.A.A',
+        'C',
+      ]);
+      act(() => apiRef.current?.collapseAllRows());
+      expect(getColumnValues(1)).to.deep.equal(['A', 'B', 'C']);
+    });
+  });
+
   describe('prop: groupingColDef', () => {
     it('should set the custom headerName', () => {
       render(<Test groupingColDef={{ headerName: 'Custom header name' }} />);

packages/x-data-grid-pro/src/utils/tree/insertDataRowInTree.ts

@@ -118,6 +118,7 @@ export const insertDataRowInTree = ({
             node,
             defaultGroupingExpansionDepth,
             isGroupExpandedByDefault,
+            (previousTree?.[id] as GridGroupNode)?.childrenExpanded,
           );
           if (shouldFetchChildren) {
             groupsToFetch?.add(id);
@@ -192,6 +193,7 @@ export const insertDataRowInTree = ({
           autoGeneratedGroupNode,
           defaultGroupingExpansionDepth,
           isGroupExpandedByDefault,
+          (previousTree?.[nodeId] as GridGroupNode)?.childrenExpanded,
         ),
         tree,
         treeDepths,
@@ -223,6 +225,7 @@ export const insertDataRowInTree = ({
           groupNode,
           defaultGroupingExpansionDepth,
           isGroupExpandedByDefault,
+          (previousTree?.[currentGroupNode.id] as GridGroupNode)?.childrenExpanded,
         );
       }
       parentNodeId = currentGroupNode.id;

packages/x-data-grid-pro/src/utils/tree/utils.ts

@@ -54,6 +54,7 @@ export const checkGroupChildrenExpansion = (
   node: GridGroupNode,
   defaultGroupingExpansionDepth: number,
   isGroupExpandedByDefault?: DataGridProProps['isGroupExpandedByDefault'],
+  prevChildrenExpanded?: boolean,
 ) => {
   let childrenExpanded: boolean;
   if (node.id === GRID_ROOT_GROUP_ID) {
@@ -62,7 +63,9 @@ export const checkGroupChildrenExpansion = (
     childrenExpanded = isGroupExpandedByDefault(node);
   } else {
     childrenExpanded =
-      defaultGroupingExpansionDepth === -1 || defaultGroupingExpansionDepth > node.depth;
+      defaultGroupingExpansionDepth === -1 ||
+      defaultGroupingExpansionDepth > node.depth ||
+      (prevChildrenExpanded ?? false);
   }
 
   return childrenExpanded;
@@ -72,11 +75,13 @@ export const updateGroupDefaultExpansion = (
   node: GridGroupNode,
   defaultGroupingExpansionDepth: number,
   isGroupExpandedByDefault?: DataGridProProps['isGroupExpandedByDefault'],
+  prevChildrenExpanded?: boolean,
 ) => {
   const childrenExpanded = checkGroupChildrenExpansion(
     node,
     defaultGroupingExpansionDepth,
     isGroupExpandedByDefault,
+    prevChildrenExpanded,
   );
   node.childrenExpanded = childrenExpanded;
   return node;

packages/x-data-grid/src/hooks/features/rows/useGridRows.ts

@@ -302,6 +302,45 @@ export const useGridRows = (
     [apiRef],
   );
 
+  const expandAllRows = React.useCallback<GridRowProApi['expandAllRows']>(() => {
+    const tree = { ...gridRowTreeSelector(apiRef) };
+
+    const traverse = (nodeId: GridRowId) => {
+      const node = tree[nodeId];
+      if (node?.type === 'group') {
+        tree[nodeId] = { ...node, childrenExpanded: true };
+        node.children.forEach(traverse);
+      }
+    };
+    traverse(GRID_ROOT_GROUP_ID);
+
+    apiRef.current.setState((state) => ({
+      ...state,
+      rows: { ...state.rows, tree },
+    }));
+    apiRef.current.publishEvent('rowExpansionChange', tree[GRID_ROOT_GROUP_ID] as GridGroupNode);
+  }, [apiRef]);
+
+  const collapseAllRows = React.useCallback<GridRowProApi['collapseAllRows']>(() => {
+    const tree = { ...gridRowTreeSelector(apiRef) };
+
+    const traverse = (nodeId: GridRowId) => {
+      const node = tree[nodeId];
+      if (node?.type === 'group') {
+        tree[nodeId] = { ...node, childrenExpanded: false };
+        node.children.forEach(traverse);
+      }
+    };
+    traverse(GRID_ROOT_GROUP_ID);
+
+    apiRef.current.setState((state) => ({
+      ...state,
+      rows: { ...state.rows, tree },
+    }));
+
+    apiRef.current.publishEvent('rowExpansionChange', tree[GRID_ROOT_GROUP_ID] as GridGroupNode);
+  }, [apiRef]);
+
   const getRowNode = React.useCallback<GridRowApi['getRowNode']>(
     (id) => (gridRowNodeSelector(apiRef, id) as any) ?? null,
     [apiRef],
@@ -455,6 +494,8 @@ export const useGridRows = (
     setRowIndex,
     setRowChildrenExpansion,
     getRowGroupChildren,
+    expandAllRows,
+    collapseAllRows,
   };
 
   const rowProPrivateApi: GridRowProPrivateApi = {

packages/x-data-grid/src/models/api/gridRowApi.ts

@@ -118,10 +118,20 @@ export interface GridRowProApi {
   getRowGroupChildren: (params: GridRowGroupChildrenGetterParams) => GridRowId[];
   /**
    * Expand or collapse a row children.
+   * Only works for the client side data or to collapse already fetched server side rows.
+   * For server-side data, use `dataSource.fetchRows(childId)` to fetch and expand the children.
    * @param {GridRowId} id the ID of the row to expand or collapse.
    * @param {boolean} isExpanded A boolean indicating if the row must be expanded or collapsed.
    */
   setRowChildrenExpansion: (id: GridRowId, isExpanded: boolean) => void;
+  /**
+   * Expand all rows. Works for the client side data only.
+   */
+  expandAllRows: () => void;
+  /**
+   * Collapse all rows. Works for the client side data only.
+   */
+  collapseAllRows: () => void;
 }
 
 export interface GridRowProPrivateApi {