DataGridPro Modify TreeGrid node icon based on data

Question

I am using MUI DataGridPro and I want to modify TreeGrid node icon based on data. This is what I want:

But from what I have studied from documentation and from what I have tried, I am only able to show static set of folders using slots:

<DataGridPro
  treeData
  apiRef={dataGridApi}
  ...
  slots={{
      treeDataCollapseIcon: () => <HierarchyFolderIcon collapsed={true} /> ,
      treeDataExpandIcon: () => <HierarchyFolderIcon collapsed={false} />   
    }} 
/>

How can I modify these icons based on the row content? slots do not provide a row or params prop.

Answer 1

This is how you can do it -

import * as React from "react";
import {
  DataGridPro,
  GridRenderCellParams,
  useGridApiContext,
  GridColDef,
  GridRowsProp,
  DataGridProProps,
  useGridSelector,
  GridKeyValue,
  gridFilteredDescendantCountLookupSelector
} from "@mui/x-data-grid-pro";
import Box, { BoxProps } from "@mui/material/Box";
import {
  DriveFolderUpload,
  FolderShared,
  FolderDelete,
  Folder,
  InsertPhoto,
  VideoFile,
  PictureAsPdf,
  InsertDriveFile,
  Description as Document,
  ExpandMore,
  KeyboardArrowRight as ExpandLess
} from "@mui/icons-material";
import { Typography } from "@mui/material";

declare module "@mui/x-data-grid/models/gridRows" {
  export interface GridLeafNode {
    gType: string;
    groupingKey: GridKeyValue | null;
    childrenExpanded: boolean;
  }

  export interface GridDataGroupNode {
    gType: string;
    groupingKey: string;
  }

  export interface GridAutoGeneratedGroupNode {
    gType: string;
    groupingKey: string;
  }

  export interface GridFooterNode {
    gType: string;
    groupingKey: string;
    childrenExpanded: boolean;
  }

  export interface GridDataPinnedRowNode {
    gType: string;
    groupingKey: string;
    childrenExpanded: boolean;
  }

  export interface GridAutoGeneratedPinnedRowNode {
    gType: string;
    groupingKey: string;
    childrenExpanded: boolean;
  }
}

type StringWithAutoComplete<T extends string> =
  | T
  | (string & Record<never, never>);

type GroupTypes = StringWithAutoComplete<
  "dir" | "file" | "img" | "video" | "doc" | "pdf"
>;

function CustomGridTreeDataGroupingCell(props: GridRenderCellParams) {
  const { id, field, rowNode } = props;
  const [cellUpdated, setCellUpdated] = React.useState(false);
  const apiRef = useGridApiContext();
  const filteredDescendantCountLookup = useGridSelector(
    apiRef,
    gridFilteredDescendantCountLookupSelector
  );
  const filteredDescendantCount =
    filteredDescendantCountLookup[rowNode.id] ?? 0;

  const handleClick: BoxProps["onClick"] = (event) => {
    if (rowNode.type !== "group") {
      return;
    }

    apiRef.current.setRowChildrenExpansion(id, !rowNode.childrenExpanded);
    apiRef.current.setCellFocus(id, field);
    event.stopPropagation();
  };

  // We need to re-render the cells to get the data they need
  // to show the icons, otherwise a default icon of
  // <InsertDriveFile /> is shown, this forces them to update it.
  React.useEffect(() => {
    if (!cellUpdated) {
      setCellUpdated(true);
    }
  }, [cellUpdated]);

  const Icon = (props: { type: string }) => {
    switch (props.type) {
      case "drive":
        return <DriveFolderUpload />;
      case "shared":
        return <FolderShared />;
      case "trash":
        return <FolderDelete />;
      case "dir":
        return <Folder />;
      case "folder":
        return <Folder />;
      case "img":
        return <InsertPhoto />;
      case "video":
        return <VideoFile />;
      case "pdf":
        return <PictureAsPdf />;
      case "file":
        return <InsertDriveFile />;
      case "doc":
        return <Document />;
      default:
        return <InsertDriveFile />;
    }
  };

  return (
    <Box sx={{ ml: rowNode.depth * 4 }}>
      <div>
        {filteredDescendantCount > 0 ? (
          <Box
            onClick={handleClick}
            tabIndex={-1}
            sx={{ display: "flex", alignItems: "center", gap: 0.5 }}
          >
            {rowNode.childrenExpanded ? <ExpandMore /> : <ExpandLess />}
            <Icon type={rowNode.gType} />
            <Typography
              sx={{ color: "#1976d2", cursor: "pointer", fontSize: "14px" }}
            >
              {rowNode.groupingKey} ({filteredDescendantCount})
            </Typography>
          </Box>
        ) : (
          <Icon type={rowNode.gType} />
        )}
      </div>
    </Box>
  );
}

interface Row {
  hierarchy: string[];
  dateModified: Date;
  id: number;
  gType: GroupTypes;
  icon?: string;
}

const rows: GridRowsProp<Row> = [
  {
    hierarchy: ["Drive"],
    dateModified: new Date(2017, 3, 4),
    id: 1,
    gType: "dir",
    icon: "drive"
  },
  {
    hierarchy: ["Drive", "Screenshot"],
    dateModified: new Date(2020, 11, 20),
    id: 2,
    gType: "file"
  },
  {
    hierarchy: ["Drive", "image01.png"],
    dateModified: new Date(2020, 10, 14),
    id: 3,
    gType: "img"
  },
  {
    hierarchy: ["Drive", "shot.mp4"],
    dateModified: new Date(2017, 10, 29),
    id: 4,
    gType: "video"
  },
  {
    hierarchy: ["Drive", "imp_docs.docx"],
    dateModified: new Date(2020, 7, 21),
    id: 5,
    gType: "doc"
  },
  {
    hierarchy: ["Drive", "read.pdf"],
    dateModified: new Date(2020, 7, 20),
    id: 6,
    gType: "pdf"
  },
  {
    hierarchy: ["Drive", "archive.zip"],
    dateModified: new Date(2019, 6, 28),
    id: 7,
    gType: "file"
  },
  {
    hierarchy: ["Trash"],
    dateModified: new Date(2016, 3, 14),
    id: 8,
    gType: "dir",
    icon: "trash"
  },
  {
    hierarchy: ["Trash", "Secret Folder"],
    dateModified: new Date(2016, 5, 17),
    id: 9,
    gType: "dir",
    icon: "folder"
  },
  {
    hierarchy: ["Trash", "Secret Folder", "barely_legal.jpg"],
    dateModified: new Date(2019, 11, 7),
    id: 10,
    gType: "img"
  },
  {
    hierarchy: ["Trash", "video_4453.mp4"],
    dateModified: new Date(2021, 7, 1),
    id: 11,
    gType: "video"
  },
  {
    hierarchy: ["Trash", "Shared"],
    dateModified: new Date(2017, 0, 12),
    id: 12,
    gType: "dir",
    icon: "shared"
  },
  {
    hierarchy: ["Trash", "Shared", "uknown"],
    dateModified: new Date(2019, 2, 22),
    id: 13,
    gType: "file"
  },
  {
    hierarchy: ["Trash", "Shared", "uknown2"],
    dateModified: new Date(2018, 4, 19),
    id: 14,
    gType: "file"
  }
];

const columns: GridColDef[] = [
  {
    field: "name",
    headerName: " External Name",
    width: 200,
    valueGetter: (params) => {
      params.rowNode.gType = params.row.icon || params.row.gType;
      const hierarchy = params.row.hierarchy;
      return hierarchy[hierarchy.length - 1];
    }
  } as GridColDef<Row, string>,
  {
    field: "dateModified",
    headerName: "Date Modified",
    width: 150,
    type: "date"
  },
  {
    field: "contentType",
    headerName: "Content Type",
    width: 150,
    type: "string",
    valueGetter: (params) => {
      switch (params.row.gType) {
        case "dir":
          return "Directory";
        case "file":
          return "File";
        case "img":
          return "Image";
        case "video":
          return "Video";
        case "doc":
          return "Document";
        case "pdf":
          return "PDF";
        default:
          return "Unknown";
      }
    }
  }
];

const getTreeDataPath: DataGridProProps["getTreeDataPath"] = (row) =>
  row.hierarchy;

const groupingColDef: DataGridProProps["groupingColDef"] = {
  headerName: "Content",
  renderCell: (params) => <CustomGridTreeDataGroupingCell {...params} />
};

export default function TreeDataCustomGroupingColumn() {
  return (
    <div style={{ height: 400, width: "100%" }}>
      <DataGridPro
        treeData
        rows={rows}
        columns={columns}
        getTreeDataPath={getTreeDataPath}
        groupingColDef={groupingColDef}
      />
    </div>
  );
}

This is what it looks like -

Here's Codesandbox link.

Mind that this is in Typescript, for Javascript, just remove all the typings.