feat: TreeSelect support maxCount

docs/demo/mutiple-with-maxCount.md

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

examples/mutiple-with-maxCount.tsx

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

src/OptionList.tsx

@@ -9,7 +9,7 @@ import * as React from 'react';
 import LegacyContext from './LegacyContext';
 import TreeSelectContext from './TreeSelectContext';
 import type { Key, SafeKey } from './interface';
-import { getAllKeys, isCheckDisabled } from './utils/valueUtil';
+import { getAllKeys, isCheckDisabled, isValidCount } from './utils/valueUtil';
 
 const HIDDEN_STYLE = {
   width: 0,
@@ -45,6 +45,7 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     treeExpandAction,
     treeTitleRender,
     onPopupScroll,
+    maxCount,
   } = React.useContext(TreeSelectContext);
 
   const {
@@ -76,6 +77,24 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     (prev, next) => next[0] && prev[1] !== next[1],
   );
 
+  const isOverMaxCount = React.useMemo<boolean>(
+    () => multiple && isValidCount(maxCount) && checkedKeys.length >= maxCount,
+    [checkedKeys, maxCount, multiple],
+  );
+
+  const traverse = (nodes: EventDataNode<any>[]): EventDataNode<any>[] => {
+    return nodes.map(node => ({
+      ...node,
+      disabled: isOverMaxCount && !checkedKeys.includes(node.key as SafeKey) ? true : node.disabled,
+      children: node.children ? traverse(node.children) : undefined,
+    }));
+  };
+
+  const processedTreeData = React.useMemo(
+    () => traverse(memoTreeData),
+    [memoTreeData, isOverMaxCount, checkedKeys],
+  );
+
   // ========================== Values ==========================
   const mergedCheckedKeys = React.useMemo(() => {
     if (!checkable) {
@@ -109,8 +128,10 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
       return;
     }
 
+    const isSelected = !checkedKeys.includes(node.key);
+
     onSelect(node.key, {
-      selected: !checkedKeys.includes(node.key),
+      selected: isSelected,
     });
 
     if (!multiple) {
@@ -183,6 +204,13 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
   const [activeKey, setActiveKey] = React.useState<Key>(null);
   const activeEntity = keyEntities[activeKey as SafeKey];
 
+  const onActiveChange = (key: Key) => {
+    if (isOverMaxCount && !checkedKeys.includes(key)) {
+      return;
+    }
+    setActiveKey(key);
+  };
+
   React.useEffect(() => {
     if (!open) {
       return;
@@ -279,7 +307,7 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
         ref={treeRef}
         focusable={false}
         prefixCls={`${prefixCls}-tree`}
-        treeData={memoTreeData}
+        treeData={processedTreeData}
         height={listHeight}
         itemHeight={listItemHeight}
         itemScrollOffset={listItemScrollOffset}
@@ -301,7 +329,7 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
         titleRender={treeTitleRender}
         {...treeProps}
         // Proxy event out
-        onActiveChange={setActiveKey}
+        onActiveChange={onActiveChange}
         onSelect={onInternalSelect}
         onCheck={onInternalSelect}
         onExpand={onInternalExpand}

src/TreeSelect.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,
@@ -558,6 +560,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       onDeselect,
       rawCheckedValues,
       rawHalfCheckedValues,
+      maxCount,
     ],
   );
 
@@ -609,6 +612,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeExpandAction,
       treeTitleRender,
       onPopupScroll,
+      maxCount,
     }),
     [
       virtual,
@@ -622,6 +626,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeExpandAction,
       treeTitleRender,
       onPopupScroll,
+      maxCount,
     ],
   );
 

src/TreeSelectContext.ts

@@ -14,6 +14,7 @@ export interface TreeSelectContextProps {
   treeExpandAction?: ExpandAction;
   treeTitleRender?: (node: any) => React.ReactNode;
   onPopupScroll?: React.UIEventHandler<HTMLDivElement>;
+  maxCount?: number;
 }
 
 const TreeSelectContext = React.createContext<TreeSelectContextProps>(null as any);

src/utils/valueUtil.ts

@@ -35,3 +35,6 @@ export const getAllKeys = (treeData: DataNode[], fieldNames: FieldNames): SafeKe
 };
 
 export const isNil = (val: any): boolean => val === null || val === undefined;
+
+export const isValidCount = (value?: number) =>
+  typeof value !== 'undefined' && !Number.isNaN(value);

tests/Select.SearchInput.spec.js

@@ -221,14 +221,22 @@ describe('TreeSelect.SearchInput', () => {
       const input = getByRole('combobox');
       fireEvent.change(input, { target: { value: '1' } });
       fireEvent.keyDown(input, { keyCode: KeyCode.ENTER });
+      fireEvent.keyUp(input, { keyCode: KeyCode.ENTER });
       expect(onSelect).toHaveBeenCalledWith('1', expect.anything());
       onSelect.mockReset();
 
       // Search disabled node and press enter, should not select
       fireEvent.change(input, { target: { value: '2' } });
       fireEvent.keyDown(input, { keyCode: KeyCode.ENTER });
+      fireEvent.keyUp(input, { keyCode: KeyCode.ENTER });
       expect(onSelect).not.toHaveBeenCalled();
       onSelect.mockReset();
+
+      // Search and press enter, should select first matched non-disabled node
+      fireEvent.change(input, { target: { value: '3' } });
+      fireEvent.keyDown(input, { keyCode: KeyCode.ENTER });
+      fireEvent.keyUp(input, { keyCode: KeyCode.ENTER });
+      expect(onSelect).toHaveBeenCalledWith('3', expect.anything());
     });
 
     it('should not select node when no matches found', () => {