[TreeView] New instance and publicAPI method: getItem

packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.types.ts

@@ -11,9 +11,8 @@ export interface UseTreeViewFocusInstance {
   focusDefaultNode: (event: React.SyntheticEvent) => void;
   focusRoot: () => void;
 }
-export interface UseTreeViewFocusPublicAPI {
-  focusNode: (event: React.SyntheticEvent, nodeId: string | null) => void;
-}
+
+export interface UseTreeViewFocusPublicAPI extends Pick<UseTreeViewFocusInstance, 'focusNode'> {}
 
 export interface UseTreeViewFocusParameters {
   /**

packages/x-tree-view/src/internals/plugins/useTreeViewNodes/useTreeViewNodes.ts

@@ -1,27 +1,29 @@
 import * as React from 'react';
 import useEventCallback from '@mui/utils/useEventCallback';
 import { TreeViewPlugin } from '../../models';
-import { populateInstance } from '../../useTreeView/useTreeView.utils';
+import { populateInstance, populatePublicAPI } from '../../useTreeView/useTreeView.utils';
 import {
   UseTreeViewNodesSignature,
   UseTreeViewNodesDefaultizedParameters,
   TreeViewNodeMap,
   TreeViewNodeIdAndChildren,
   UseTreeViewNodesState,
+  TreeViewItemMap,
 } from './useTreeViewNodes.types';
 import { publishTreeViewEvent } from '../../utils/publishTreeViewEvent';
 import { TreeViewBaseItem } from '../../../models';
 
-const updateState = ({
+const updateNodesState = ({
   items,
   isItemDisabled,
   getItemLabel,
   getItemId,
 }: Pick<
   UseTreeViewNodesDefaultizedParameters<TreeViewBaseItem>,
   'items' | 'isItemDisabled' | 'getItemLabel' | 'getItemId'
->): UseTreeViewNodesState => {
+>): UseTreeViewNodesState<any>['nodes'] => {
   const nodeMap: TreeViewNodeMap = {};
+  const itemMap: TreeViewItemMap<any> = {};
 
   const processItem = (
     item: TreeViewBaseItem,
@@ -73,6 +75,8 @@ const updateState = ({
       disabled: isItemDisabled ? isItemDisabled(item) : false,
     };
 
+    itemMap[id] = item;
+
     return {
       id,
       children: item.children?.map((child, childIndex) => processItem(child, childIndex, id)),
@@ -84,6 +88,7 @@ const updateState = ({
   return {
     nodeMap,
     nodeTree,
+    itemMap,
   };
 };
 
@@ -93,7 +98,15 @@ export const useTreeViewNodes: TreeViewPlugin<UseTreeViewNodesSignature> = ({
   state,
   setState,
 }) => {
-  const getNode = React.useCallback((nodeId: string) => state.nodeMap[nodeId], [state.nodeMap]);
+  const getNode = React.useCallback(
+    (nodeId: string) => state.nodes.nodeMap[nodeId],
+    [state.nodes.nodeMap],
+  );
+
+  const getItem = React.useCallback(
+    (nodeId: string) => state.nodes.itemMap[nodeId],
+    [state.nodes.itemMap],
+  );
 
   const isNodeDisabled = React.useCallback(
     (nodeId: string | null): nodeId is string => {
@@ -125,7 +138,7 @@ export const useTreeViewNodes: TreeViewPlugin<UseTreeViewNodesSignature> = ({
   );
 
   const getChildrenIds = useEventCallback((nodeId: string | null) =>
-    Object.values(state.nodeMap)
+    Object.values(state.nodes.nodeMap)
       .filter((node) => node.parentId === nodeId)
       .sort((a, b) => a.index - b.index)
       .map((child) => child.id),
@@ -142,14 +155,14 @@ export const useTreeViewNodes: TreeViewPlugin<UseTreeViewNodesSignature> = ({
 
   React.useEffect(() => {
     setState((prevState) => {
-      const newState = updateState({
+      const newState = updateNodesState({
         items: params.items,
         isItemDisabled: params.isItemDisabled,
         getItemId: params.getItemId,
         getItemLabel: params.getItemLabel,
       });
 
-      Object.values(prevState.nodeMap).forEach((node) => {
+      Object.values(prevState.nodes.nodeMap).forEach((node) => {
         if (!newState.nodeMap[node.id]) {
           publishTreeViewEvent(instance, 'removeNode', { id: node.id });
         }
@@ -171,7 +184,7 @@ export const useTreeViewNodes: TreeViewPlugin<UseTreeViewNodesSignature> = ({
       id,
       children,
     }: TreeViewNodeIdAndChildren): ReturnType<typeof instance.getNodesToRender>[number] => {
-      const node = state.nodeMap[id];
+      const node = state.nodes.nodeMap[id];
       return {
         label: node.label!,
         nodeId: node.id,
@@ -180,29 +193,35 @@ export const useTreeViewNodes: TreeViewPlugin<UseTreeViewNodesSignature> = ({
       };
     };
 
-    return state.nodeTree.map(getPropsFromNodeId);
+    return state.nodes.nodeTree.map(getPropsFromNodeId);
   });
 
   populateInstance<UseTreeViewNodesSignature>(instance, {
     getNode,
+    getItem,
     getNodesToRender,
     getChildrenIds,
     getNavigableChildrenIds,
     isNodeDisabled,
   });
 
+  populatePublicAPI<UseTreeViewNodesSignature>(instance, {
+    getItem,
+  });
+
   return {
     contextValue: { disabledItemsFocusable: params.disabledItemsFocusable },
   };
 };
 
-useTreeViewNodes.getInitialState = (params) =>
-  updateState({
+useTreeViewNodes.getInitialState = (params) => ({
+  nodes: updateNodesState({
     items: params.items,
     isItemDisabled: params.isItemDisabled,
     getItemId: params.getItemId,
     getItemLabel: params.getItemLabel,
-  });
+  }),
+});
 
 useTreeViewNodes.getDefaultizedParams = (params) => ({
   ...params,

packages/x-tree-view/src/internals/plugins/useTreeViewNodes/useTreeViewNodes.types.ts

@@ -8,14 +8,18 @@ interface TreeViewNodeProps {
   children?: TreeViewNodeProps[];
 }
 
-export interface UseTreeViewNodesInstance {
+export interface UseTreeViewNodesInstance<R extends {}> {
   getNode: (nodeId: string) => TreeViewNode;
+  getItem: (nodeId: string) => R;
   getNodesToRender: () => TreeViewNodeProps[];
   getChildrenIds: (nodeId: string | null) => string[];
   getNavigableChildrenIds: (nodeId: string | null) => string[];
   isNodeDisabled: (nodeId: string | null) => nodeId is string;
 }
 
+export interface UseTreeViewNodesPublicAPI<R extends {}>
+  extends Pick<UseTreeViewNodesInstance<R>, 'getItem'> {}
+
 export interface UseTreeViewNodesParameters<R extends {}> {
   /**
    * If `true`, will allow focus on disabled items.
@@ -66,9 +70,12 @@ export interface TreeViewNodeIdAndChildren {
   children?: TreeViewNodeIdAndChildren[];
 }
 
-export interface UseTreeViewNodesState {
-  nodeTree: TreeViewNodeIdAndChildren[];
-  nodeMap: TreeViewNodeMap;
+export interface UseTreeViewNodesState<R extends {}> {
+  nodes: {
+    nodeTree: TreeViewNodeIdAndChildren[];
+    nodeMap: TreeViewNodeMap;
+    itemMap: TreeViewItemMap<R>;
+  };
 }
 
 interface UseTreeViewNodesContextValue
@@ -77,10 +84,12 @@ interface UseTreeViewNodesContextValue
 export type UseTreeViewNodesSignature = TreeViewPluginSignature<{
   params: UseTreeViewNodesParameters<any>;
   defaultizedParams: UseTreeViewNodesDefaultizedParameters<any>;
-  instance: UseTreeViewNodesInstance;
+  instance: UseTreeViewNodesInstance<any>;
   events: UseTreeViewNodesEventLookup;
-  state: UseTreeViewNodesState;
+  state: UseTreeViewNodesState<any>;
   contextValue: UseTreeViewNodesContextValue;
 }>;
 
 export type TreeViewNodeMap = { [nodeId: string]: TreeViewNode };
+
+export type TreeViewItemMap<R extends {}> = { [nodeId: string]: R };