[tree view] Explore a better plugin model API (#11567)

mui · Jan 12, 2024 · 69f1ca4 · 69f1ca4
1 parent d537fba
commit 69f1ca4
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 39 deletions.
diff --git a/packages/x-tree-view/src/internals/models/plugin.ts b/packages/x-tree-view/src/internals/models/plugin.ts
@@ -17,8 +17,9 @@ export interface TreeViewPluginOptions<TSignature extends TreeViewAnyPluginSigna
 
 type TreeViewModelsInitializer<TSignature extends TreeViewAnyPluginSignature> = {
   [TControlled in keyof TSignature['models']]: {
-    controlledProp: TControlled;
-    defaultProp: keyof TSignature['params'];
+    getDefaultValue: (
+      params: TSignature['defaultizedParams'],
+    ) => Exclude<TSignature['defaultizedParams'][TControlled], undefined>;
   };
 };
 
@@ -91,8 +92,13 @@ export type TreeViewUsedInstance<TSignature extends TreeViewAnyPluginSignature>
 type TreeViewUsedState<TSignature extends TreeViewAnyPluginSignature> = TSignature['state'] &
   MergePluginsProperty<TreeViewUsedPlugins<TSignature>, 'state'>;
 
+type RemoveSetValue<Models extends Record<string, TreeViewModel<any>>> = {
+  [K in keyof Models]: Omit<Models[K], 'setValue'>;
+};
+
 export type TreeViewUsedModels<TSignature extends TreeViewAnyPluginSignature> =
-  TSignature['models'] & MergePluginsProperty<TreeViewUsedPlugins<TSignature>, 'models'>;
+  TSignature['models'] &
+    RemoveSetValue<MergePluginsProperty<TreeViewUsedPlugins<TSignature>, 'models'>>;
 
 export type TreeViewUsedEvents<TSignature extends TreeViewAnyPluginSignature> =
   TSignature['events'] & MergePluginsProperty<TreeViewUsedPlugins<TSignature>, 'events'>;

diff --git a/packages/x-tree-view/src/internals/models/treeView.ts b/packages/x-tree-view/src/internals/models/treeView.ts
@@ -1,4 +1,3 @@
-import * as React from 'react';
 import type { TreeViewAnyPluginSignature } from './plugin';
 import type { MergePluginsProperty } from './helpers';
 
@@ -25,7 +24,7 @@ export interface TreeViewItemRange {
 export interface TreeViewModel<TValue> {
   name: string;
   value: TValue;
-  setValue: React.Dispatch<React.SetStateAction<TValue>>;
+  setControlledValue: (value: TValue | ((prevValue: TValue) => TValue)) => void;
 }
 
 export type TreeViewInstance<TSignatures extends readonly TreeViewAnyPluginSignature[]> =

diff --git a/packages/x-tree-view/src/internals/plugins/useTreeViewExpansion/useTreeViewExpansion.ts b/packages/x-tree-view/src/internals/plugins/useTreeViewExpansion/useTreeViewExpansion.ts
@@ -9,6 +9,11 @@ export const useTreeViewExpansion: TreeViewPlugin<UseTreeViewExpansionSignature>
   params,
   models,
 }) => {
+  const setExpandedNodes = (event: React.SyntheticEvent, value: string[]) => {
+    params.onExpandedNodesChange?.(event, value);
+    models.expandedNodes.setControlledValue(value);
+  };
+
   const isNodeExpanded = React.useCallback(
     (nodeId: string) => {
       return Array.isArray(models.expandedNodes.value)
@@ -42,11 +47,7 @@ export const useTreeViewExpansion: TreeViewPlugin<UseTreeViewExpansionSignature>
         params.onNodeExpansionToggle(event, nodeId, !isExpandedBefore);
       }
 
-      if (params.onExpandedNodesChange) {
-        params.onExpandedNodesChange(event, newExpanded);
-      }
-
-      models.expandedNodes.setValue(newExpanded);
+      setExpandedNodes(event, newExpanded);
     },
   );
 
@@ -67,11 +68,7 @@ export const useTreeViewExpansion: TreeViewPlugin<UseTreeViewExpansionSignature>
         });
       }
 
-      if (params.onExpandedNodesChange) {
-        params.onExpandedNodesChange(event, newExpanded);
-      }
-
-      models.expandedNodes.setValue(newExpanded);
+      setExpandedNodes(event, newExpanded);
     }
   };
 
@@ -85,8 +82,7 @@ export const useTreeViewExpansion: TreeViewPlugin<UseTreeViewExpansionSignature>
 
 useTreeViewExpansion.models = {
   expandedNodes: {
-    controlledProp: 'expandedNodes',
-    defaultProp: 'defaultExpandedNodes',
+    getDefaultValue: (params) => params.defaultExpandedNodes,
   },
 };
 

diff --git a/packages/x-tree-view/src/internals/plugins/useTreeViewSelection/useTreeViewSelection.ts b/packages/x-tree-view/src/internals/plugins/useTreeViewSelection/useTreeViewSelection.ts
@@ -18,11 +18,6 @@ export const useTreeViewSelection: TreeViewPlugin<UseTreeViewSelectionSignature>
   const lastSelectionWasRange = React.useRef(false);
   const currentRangeSelection = React.useRef<string[]>([]);
 
-  const isNodeSelected = (nodeId: string) =>
-    Array.isArray(models.selectedNodes.value)
-      ? models.selectedNodes.value.indexOf(nodeId) !== -1
-      : models.selectedNodes.value === nodeId;
-
   const setSelectedNodes = (
     event: React.SyntheticEvent,
     newSelectedNodes: typeof params.defaultSelectedNodes,
@@ -57,9 +52,14 @@ export const useTreeViewSelection: TreeViewPlugin<UseTreeViewSelectionSignature>
       params.onSelectedNodesChange(event, newSelectedNodes);
     }
 
-    models.selectedNodes.setValue(newSelectedNodes);
+    models.selectedNodes.setControlledValue(newSelectedNodes);
   };
 
+  const isNodeSelected = (nodeId: string) =>
+    Array.isArray(models.selectedNodes.value)
+      ? models.selectedNodes.value.indexOf(nodeId) !== -1
+      : models.selectedNodes.value === nodeId;
+
   const selectNode = (event: React.SyntheticEvent, nodeId: string, multiple = false) => {
     if (params.disableSelection) {
       return;
@@ -125,7 +125,6 @@ export const useTreeViewSelection: TreeViewPlugin<UseTreeViewSelectionSignature>
       base.push(next);
       currentRangeSelection.current.push(current, next);
     }
-
     setSelectedNodes(event, base);
   };
 
@@ -145,7 +144,6 @@ export const useTreeViewSelection: TreeViewPlugin<UseTreeViewSelectionSignature>
     currentRangeSelection.current = range;
     let newSelected = base.concat(range);
     newSelected = newSelected.filter((id, i) => newSelected.indexOf(id) === i);
-
     setSelectedNodes(event, newSelected);
   };
 
@@ -210,7 +208,9 @@ export const useTreeViewSelection: TreeViewPlugin<UseTreeViewSelectionSignature>
 };
 
 useTreeViewSelection.models = {
-  selectedNodes: { controlledProp: 'selectedNodes', defaultProp: 'defaultSelectedNodes' },
+  selectedNodes: {
+    getDefaultValue: (params) => params.defaultSelectedNodes,
+  },
 };
 
 const DEFAULT_SELECTED_NODES: string[] = [];

diff --git a/packages/x-tree-view/src/internals/useTreeView/useTreeViewModels.ts b/packages/x-tree-view/src/internals/useTreeView/useTreeViewModels.ts
@@ -16,10 +16,14 @@ export const useTreeViewModels = <
   plugins: TPlugins,
   props: MergePluginsProperty<ConvertPluginsIntoSignatures<TPlugins>, 'defaultizedParams'>,
 ) => {
+  type DefaultizedParams = MergePluginsProperty<
+    ConvertPluginsIntoSignatures<TPlugins>,
+    'defaultizedParams'
+  >;
+
   const modelsRef = React.useRef<{
     [modelName: string]: {
-      controlledProp: keyof typeof props;
-      defaultProp: keyof typeof props;
+      getDefaultValue: (params: DefaultizedParams) => any;
       isControlled: boolean;
     };
   }>({});
@@ -29,13 +33,12 @@ export const useTreeViewModels = <
 
     plugins.forEach((plugin) => {
       if (plugin.models) {
-        Object.entries(plugin.models).forEach(([modelName, model]) => {
+        Object.entries(plugin.models).forEach(([modelName, modelInitializer]) => {
           modelsRef.current[modelName] = {
-            controlledProp: model.controlledProp as keyof typeof props,
-            defaultProp: model.defaultProp as keyof typeof props,
-            isControlled: props[model.controlledProp as keyof typeof props] !== undefined,
+            isControlled: props[modelName as keyof DefaultizedParams] !== undefined,
+            getDefaultValue: modelInitializer.getDefaultValue,
           };
-          initialState[modelName] = props[model.defaultProp as keyof typeof props];
+          initialState[modelName] = modelInitializer.getDefaultValue(props);
         });
       }
     });
@@ -45,13 +48,15 @@ export const useTreeViewModels = <
 
   const models = Object.fromEntries(
     Object.entries(modelsRef.current).map(([modelName, model]) => {
-      const value = model.isControlled ? props[model.controlledProp] : modelsState[modelName];
+      const value = model.isControlled
+        ? props[modelName as keyof DefaultizedParams]
+        : modelsState[modelName];
 
       return [
         modelName,
         {
           value,
-          setValue: (newValue: any) => {
+          setControlledValue: (newValue: any) => {
             if (!model.isControlled) {
               setModelsState((prevState) => ({
                 ...prevState,
@@ -68,8 +73,8 @@ export const useTreeViewModels = <
   /* eslint-disable react-hooks/rules-of-hooks, react-hooks/exhaustive-deps */
   if (process.env.NODE_ENV !== 'production') {
     Object.entries(modelsRef.current).forEach(([modelName, model]) => {
-      const controlled = props[model.controlledProp];
-      const defaultProp = props[model.defaultProp];
+      const controlled = props[modelName as keyof DefaultizedParams];
+      const newDefaultValue = model.getDefaultValue(props);
 
       React.useEffect(() => {
         if (model.isControlled !== (controlled !== undefined)) {
@@ -90,10 +95,10 @@ export const useTreeViewModels = <
         }
       }, [controlled]);
 
-      const { current: defaultValue } = React.useRef(defaultProp);
+      const { current: defaultValue } = React.useRef(newDefaultValue);
 
       React.useEffect(() => {
-        if (!model.isControlled && defaultValue !== defaultProp) {
+        if (!model.isControlled && defaultValue !== newDefaultValue) {
           console.error(
             [
               `MUI X: A component is changing the default ${modelName} state of an uncontrolled TreeView after being initialized. ` +