feat: adjust maxCount implement logic

Author: aojunhao123

examples/mutiple-with-maxCount.tsx

@@ -21,6 +21,11 @@ export default () => {
           value: '1-2',
           title: '1-2',
         },
+        {
+          key: '1-3',
+          value: '1-3',
+          title: '1-3',
+        },
       ],
     },
     {
@@ -58,7 +63,10 @@ export default () => {
         style={{ width: 300 }}
         multiple
         treeCheckable
-        maxCount={3}
+        // showCheckedStrategy="SHOW_ALL"
+        showCheckedStrategy="SHOW_PARENT"
+        // showCheckedStrategy="SHOW_CHILD"
+        maxCount={4}
         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, isValidCount } from './utils/valueUtil';
+import { getAllKeys, isCheckDisabled } from './utils/valueUtil';
 
 const HIDDEN_STYLE = {
   width: 0,
@@ -46,6 +46,7 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     treeTitleRender,
     onPopupScroll,
     maxCount,
+    displayValues,
   } = React.useContext(TreeSelectContext);
 
   const {
@@ -77,24 +78,6 @@ 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) {
@@ -202,13 +185,6 @@ 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;
@@ -274,8 +250,36 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
       preSearchValue !== nextSearchValue && !!(nextSearchValue || nextExcludeSearchExpandedKeys),
   );
 
+  const onActiveChange = (key: Key) => {
+    if (isOverMaxCount && !displayValues?.some(v => v.value === key)) {
+      return;
+    }
+    setActiveKey(key);
+  };
+
+  const isOverMaxCount = React.useMemo<boolean>(
+    () => multiple && maxCount !== undefined && displayValues?.length >= maxCount,
+    [multiple, maxCount, displayValues?.length],
+  );
+
+  const traverse = (nodes: EventDataNode<any>[]): EventDataNode<any>[] => {
+    return nodes.map(node => ({
+      ...node,
+      disabled:
+        isOverMaxCount && !displayValues?.some(v => v.value === node[fieldNames.value])
+          ? true
+          : node.disabled,
+      children: node.children ? traverse(node.children) : undefined,
+    }));
+  };
+
+  const processedTreeData = React.useMemo(
+    () => traverse(memoTreeData),
+    [memoTreeData, isOverMaxCount, displayValues, fieldNames],
+  );
+
   // ========================== Render ==========================
-  if (memoTreeData.length === 0) {
+  if (processedTreeData.length === 0) {
     return (
       <div role="listbox" className={`${prefixCls}-empty`} onMouseDown={onListMouseDown}>
         {notFoundContent}

src/TreeSelect.tsx

@@ -415,11 +415,16 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       extra: { triggerValue?: SafeKey; selected?: boolean },
       source: SelectSource,
     ) => {
-      // if treeCheckable and not treeCheckStrictly, check if exceed maxCount
-      if (treeCheckable && !treeCheckStrictly && maxCount !== undefined) {
-        const { checkedKeys } = conductCheck(newRawValues, true, keyEntities);
+      // if multiple and maxCount is set, check if exceed maxCount
+      if (mergedMultiple && maxCount !== undefined) {
+        const formattedKeyList = formatStrategyValues(
+          newRawValues,
+          mergedShowCheckedStrategy,
+          keyEntities,
+          mergedFieldNames,
+        );
 
-        if (checkedKeys.length > maxCount) {
+        if (formattedKeyList.length > maxCount) {
           return;
         }
       }
@@ -622,6 +627,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeTitleRender,
       onPopupScroll,
       maxCount,
+      displayValues: cachedDisplayValues,
     }),
     [
       virtual,
@@ -636,6 +642,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeTitleRender,
       onPopupScroll,
       maxCount,
+      cachedDisplayValues,
     ],
   );
 

src/TreeSelectContext.ts

@@ -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;
@@ -15,6 +15,7 @@ export interface TreeSelectContextProps {
   treeTitleRender?: (node: any) => React.ReactNode;
   onPopupScroll?: React.UIEventHandler<HTMLDivElement>;
   maxCount?: number;
+  displayValues?: LabeledValueType[];
 }
 
 const TreeSelectContext = React.createContext<TreeSelectContextProps>(null as any);

src/utils/valueUtil.ts

@@ -35,6 +35,3 @@ 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);