feat: TreeSelect support maxCount (#596)

* feat: TreeSelect support maxCount

* feat: sync activeKey state

* feat: sync disabled state

* feat: sync disabled state

* docs: improve maxCount demo

* test: add maxCount test cases

* chore: remove deadCode

* test: add test cases for keyboard operations

* chore: remove useless code

* test: add test case

* test: improve test case

* docs: add maxCount description

* feat: forbid check when checkedKeys more than maxCount

* chore: demo improvement

* feat: adjust maxCount implement logic

* fix: lint fix

* test: add test cases for maxCount

* chore: hoist state to context

* chore: hoist traverse operation to TreeSelect

* feat: improve keyboard navigation when reach maxCount

* feat: improve keyboard navigation when reach maxCount

* perf: use cache to improve navigation performance

* refactor: reuse formatStrategyValues

* feat: add disabledStrategy

* feat: add code comment

* test: supplement test case for keyboard operation

* chore: handle git conflicts manually

* chore: remove useless code

* chore: memories displayValues

* refactor: use InternalContext

* chore: adjust context api

* chore: bump rc-tree version to 5.11.0 for support maxCount

* fix: test coverage

* fix: fix some case

* chore: remove keyboard operation logic

* chore: optimized code logic

Author: aojunhao123

README.md

@@ -0,0 +1,8 @@
+---
+title: mutiple-with-maxCount
+nav:
+  title: Demo
+  path: /demo
+---
+
+<code src="../../examples/mutiple-with-maxCount.tsx"></code>

docs/demo/mutiple-with-maxCount.md

@@ -0,0 +1,94 @@
+import React, { useState } from 'react';
+import TreeSelect from '../src';
+
+export default () => {
+  const [value, setValue] = useState<string[]>(['1']);
+  const [checkValue, setCheckValue] = useState<string[]>(['1']);
+
+  const treeData = [
+    {
+      key: '1',
+      value: '1',
+      title: '1',
+      children: [
+        {
+          key: '1-1',
+          value: '1-1',
+          title: '1-1',
+        },
+        {
+          key: '1-2',
+          value: '1-2',
+          title: '1-2',
+        },
+        {
+          key: '1-3',
+          value: '1-3',
+          title: '1-3',
+        },
+      ],
+    },
+    {
+      key: '2',
+      value: '2',
+      title: '2',
+    },
+    {
+      key: '3',
+      value: '3',
+      title: '3',
+    },
+    {
+      key: '4',
+      value: '4',
+      title: '4',
+    },
+  ];
+
+  const onChange = (val: string[]) => {
+    setValue(val);
+  };
+
+  const onCheckChange = (val: string[]) => {
+    setCheckValue(val);
+  };
+
+  return (
+    <>
+      <h2>multiple with maxCount</h2>
+      <TreeSelect
+        style={{ width: 300 }}
+        fieldNames={{ value: 'value', label: 'title' }}
+        multiple
+        maxCount={3}
+        treeData={treeData}
+      />
+
+      <h2>checkable with maxCount</h2>
+      <TreeSelect
+        style={{ width: 300 }}
+        multiple
+        treeCheckable
+        // showCheckedStrategy="SHOW_ALL"
+        showCheckedStrategy="SHOW_PARENT"
+        // showCheckedStrategy="SHOW_CHILD"
+        maxCount={4}
+        treeData={treeData}
+        onChange={onChange}
+        value={value}
+      />
+
+      <h2>checkable with maxCount and treeCheckStrictly</h2>
+      <TreeSelect
+        style={{ width: 300 }}
+        multiple
+        treeCheckable
+        treeCheckStrictly
+        maxCount={3}
+        treeData={treeData}
+        onChange={onCheckChange}
+        value={checkValue}
+      />
+    </>
+  );
+};

examples/mutiple-with-maxCount.tsx

@@ -47,7 +47,7 @@
     "@babel/runtime": "^7.25.7",
     "classnames": "2.x",
     "rc-select": "~14.16.2",
-    "rc-tree": "~5.10.1",
+    "rc-tree": "~5.11.0",
     "rc-util": "^5.43.0"
   },
   "devDependencies": {

package.json

@@ -2,14 +2,16 @@ import { useBaseProps } from 'rc-select';
 import type { RefOptionListProps } from 'rc-select/lib/OptionList';
 import type { TreeProps } from 'rc-tree';
 import Tree from 'rc-tree';
+import { UnstableContext } from 'rc-tree';
 import type { EventDataNode, ScrollTo } from 'rc-tree/lib/interface';
 import KeyCode from 'rc-util/lib/KeyCode';
 import useMemo from 'rc-util/lib/hooks/useMemo';
 import * as React from 'react';
 import LegacyContext from './LegacyContext';
 import TreeSelectContext from './TreeSelectContext';
-import type { Key, SafeKey } from './interface';
+import type { DataNode, Key, SafeKey } from './interface';
 import { getAllKeys, isCheckDisabled } from './utils/valueUtil';
+import { useEvent } from 'rc-util';
 
 const HIDDEN_STYLE = {
   width: 0,
@@ -45,6 +47,8 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     treeExpandAction,
     treeTitleRender,
     onPopupScroll,
+    displayValues,
+    isOverMaxCount,
   } = React.useContext(TreeSelectContext);
 
   const {
@@ -76,6 +80,11 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     (prev, next) => next[0] && prev[1] !== next[1],
   );
 
+  const memoRawValues = React.useMemo(
+    () => (displayValues || []).map(v => v.value),
+    [displayValues],
+  );
+
   // ========================== Values ==========================
   const mergedCheckedKeys = React.useMemo(() => {
     if (!checkable) {
@@ -154,6 +163,10 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [searchValue]);
 
+  const nodeDisabled = useEvent((node: DataNode) => {
+    return isOverMaxCount && !memoRawValues.includes(node[fieldNames.value]);
+  });
+
   // ========================== Get First Selectable Node ==========================
   const getFirstMatchingNode = (nodes: EventDataNode<any>[]): EventDataNode<any> | null => {
     for (const node of nodes) {
@@ -221,8 +234,9 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
         // >>> Select item
         case KeyCode.ENTER: {
           if (activeEntity) {
+            const isNodeDisabled = nodeDisabled(activeEntity.node);
             const { selectable, value, disabled } = activeEntity?.node || {};
-            if (selectable !== false && !disabled) {
+            if (selectable !== false && !disabled && !isNodeDisabled) {
               onInternalSelect(null, {
                 node: { key: activeKey },
                 selected: !checkedKeys.includes(value),
@@ -276,42 +290,43 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
           {activeEntity.node.value}
         </span>
       )}
-
-      <Tree
-        ref={treeRef}
-        focusable={false}
-        prefixCls={`${prefixCls}-tree`}
-        treeData={memoTreeData}
-        height={listHeight}
-        itemHeight={listItemHeight}
-        itemScrollOffset={listItemScrollOffset}
-        virtual={virtual !== false && dropdownMatchSelectWidth !== false}
-        multiple={multiple}
-        icon={treeIcon}
-        showIcon={showTreeIcon}
-        switcherIcon={switcherIcon}
-        showLine={treeLine}
-        loadData={syncLoadData}
-        motion={treeMotion}
-        activeKey={activeKey}
-        // We handle keys by out instead tree self
-        checkable={checkable}
-        checkStrictly
-        checkedKeys={mergedCheckedKeys}
-        selectedKeys={!checkable ? checkedKeys : []}
-        defaultExpandAll={treeDefaultExpandAll}
-        titleRender={treeTitleRender}
-        {...treeProps}
-        // Proxy event out
-        onActiveChange={setActiveKey}
-        onSelect={onInternalSelect}
-        onCheck={onInternalSelect}
-        onExpand={onInternalExpand}
-        onLoad={onTreeLoad}
-        filterTreeNode={filterTreeNode}
-        expandAction={treeExpandAction}
-        onScroll={onPopupScroll}
-      />
+      <UnstableContext.Provider value={{ nodeDisabled }}>
+        <Tree
+          ref={treeRef}
+          focusable={false}
+          prefixCls={`${prefixCls}-tree`}
+          treeData={memoTreeData}
+          height={listHeight}
+          itemHeight={listItemHeight}
+          itemScrollOffset={listItemScrollOffset}
+          virtual={virtual !== false && dropdownMatchSelectWidth !== false}
+          multiple={multiple}
+          icon={treeIcon}
+          showIcon={showTreeIcon}
+          switcherIcon={switcherIcon}
+          showLine={treeLine}
+          loadData={syncLoadData}
+          motion={treeMotion}
+          activeKey={activeKey}
+          // We handle keys by out instead tree self
+          checkable={checkable}
+          checkStrictly
+          checkedKeys={mergedCheckedKeys}
+          selectedKeys={!checkable ? checkedKeys : []}
+          defaultExpandAll={treeDefaultExpandAll}
+          titleRender={treeTitleRender}
+          {...treeProps}
+          // Proxy event out
+          onActiveChange={setActiveKey}
+          onSelect={onInternalSelect}
+          onCheck={onInternalSelect}
+          onExpand={onInternalExpand}
+          onLoad={onTreeLoad}
+          filterTreeNode={filterTreeNode}
+          expandAction={treeExpandAction}
+          onScroll={onPopupScroll}
+        />
+      </UnstableContext.Provider>
     </div>
   );
 };

src/OptionList.tsx

@@ -72,6 +72,7 @@ export interface TreeSelectProps<ValueType = any, OptionType extends DataNode =
   treeCheckable?: boolean | React.ReactNode;
   treeCheckStrictly?: boolean;
   labelInValue?: boolean;
+  maxCount?: number;
 
   // >>> Data
   treeData?: OptionType[];
@@ -136,6 +137,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
     treeCheckable,
     treeCheckStrictly,
     labelInValue,
+    maxCount,
 
     // FieldNames
     fieldNames,
@@ -413,6 +415,20 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       extra: { triggerValue?: SafeKey; selected?: boolean },
       source: SelectSource,
     ) => {
+      const formattedKeyList = formatStrategyValues(
+        newRawValues,
+        mergedShowCheckedStrategy,
+        keyEntities,
+        mergedFieldNames,
+      );
+
+      // if multiple and maxCount is set, check if exceed maxCount
+      if (mergedMultiple && maxCount !== undefined) {
+        if (formattedKeyList.length > maxCount) {
+          return;
+        }
+      }
+
       const labeledValues = convert2LabelValues(newRawValues);
       setInternalValue(labeledValues);
 
@@ -425,12 +441,6 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       if (onChange) {
         let eventValues: SafeKey[] = newRawValues;
         if (treeConduction) {
-          const formattedKeyList = formatStrategyValues(
-            newRawValues,
-            mergedShowCheckedStrategy,
-            keyEntities,
-            mergedFieldNames,
-          );
           eventValues = formattedKeyList.map(key => {
             const entity = valueEntities.get(key);
             return entity ? entity.node[mergedFieldNames.value] : key;
@@ -558,6 +568,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       onDeselect,
       rawCheckedValues,
       rawHalfCheckedValues,
+      maxCount,
     ],
   );
 
@@ -596,8 +607,11 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
   });
 
   // ========================== Context ===========================
-  const treeSelectContext = React.useMemo<TreeSelectContextProps>(
-    () => ({
+  const isOverMaxCount =
+    mergedMultiple && maxCount !== undefined && cachedDisplayValues?.length >= maxCount;
+
+  const treeSelectContext = React.useMemo<TreeSelectContextProps>(() => {
+    return {
       virtual,
       dropdownMatchSelectWidth,
       listHeight,
@@ -609,21 +623,25 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeExpandAction,
       treeTitleRender,
       onPopupScroll,
-    }),
-    [
-      virtual,
-      dropdownMatchSelectWidth,
-      listHeight,
-      listItemHeight,
-      listItemScrollOffset,
-      filteredTreeData,
-      mergedFieldNames,
-      onOptionSelect,
-      treeExpandAction,
-      treeTitleRender,
-      onPopupScroll,
-    ],
-  );
+      displayValues: cachedDisplayValues,
+      isOverMaxCount,
+    };
+  }, [
+    virtual,
+    dropdownMatchSelectWidth,
+    listHeight,
+    listItemHeight,
+    listItemScrollOffset,
+    filteredTreeData,
+    mergedFieldNames,
+    onOptionSelect,
+    treeExpandAction,
+    treeTitleRender,
+    onPopupScroll,
+    maxCount,
+    cachedDisplayValues,
+    mergedMultiple,
+  ]);
 
   // ======================= Legacy Context =======================
   const legacyContext = React.useMemo(

src/TreeSelect.tsx

@@ -1,6 +1,6 @@
 import * as React from 'react';
 import type { ExpandAction } from 'rc-tree/lib/Tree';
-import type { DataNode, FieldNames, Key } from './interface';
+import type { DataNode, FieldNames, Key, LabeledValueType } from './interface';
 
 export interface TreeSelectContextProps {
   virtual?: boolean;
@@ -14,6 +14,8 @@ export interface TreeSelectContextProps {
   treeExpandAction?: ExpandAction;
   treeTitleRender?: (node: any) => React.ReactNode;
   onPopupScroll?: React.UIEventHandler<HTMLDivElement>;
+  displayValues?: LabeledValueType[];
+  isOverMaxCount?: boolean;
 }
 
 const TreeSelectContext = React.createContext<TreeSelectContextProps>(null as any);