feat: Tree switcherIcon prop support render-prop (#34470)

* feat: Tree `switcherIcon` prop support render-prop

* add test for `switcherIcon` render-prop

* change the `switcherIcon` prop type of `TreeSelect`

* fix switcherIcon render-prop condition

* fix switcherIcon

components/tree-select/index.tsx

@@ -14,6 +14,7 @@ import { useContext } from 'react';
 import { ConfigContext } from '../config-provider';
 import devWarning from '../_util/devWarning';
 import { AntTreeNodeProps, TreeProps } from '../tree';
+import { SwitcherIcon } from '../tree/Tree';
 import getIcons from '../select/utils/iconUtil';
 import renderSwitcherIcon from '../tree/utils/iconUtil';
 import SizeContext, { SizeType } from '../config-provider/SizeContext';
@@ -50,6 +51,7 @@ export interface TreeSelectProps<
   bordered?: boolean;
   treeLine?: TreeProps['showLine'];
   status?: InputStatus;
+  switcherIcon?: SwitcherIcon;
 }
 
 const InternalTreeSelect = <OptionType extends BaseOptionType | DefaultOptionType = BaseOptionType>(

components/tree/Tree.tsx

@@ -98,6 +98,8 @@ interface DraggableConfig {
   nodeDraggable?: DraggableFn;
 }
 
+export type SwitcherIcon = ((expanded: boolean) => React.ReactNode) | React.ReactNode;
+
 export interface TreeProps<T extends BasicDataNode = DataNode>
   extends Omit<RcTreeProps<T>, 'prefixCls' | 'showLine' | 'direction' | 'draggable'> {
   showLine?: boolean | { showLeafIcon: boolean };
@@ -137,7 +139,7 @@ export interface TreeProps<T extends BasicDataNode = DataNode>
   style?: React.CSSProperties;
   showIcon?: boolean;
   icon?: ((nodeProps: AntdTreeNodeAttribute) => React.ReactNode) | React.ReactNode;
-  switcherIcon?: React.ReactElement<any>;
+  switcherIcon?: SwitcherIcon;
   prefixCls?: string;
   children?: React.ReactNode;
   blockNode?: boolean;

components/tree/__tests__/index.test.js

@@ -73,6 +73,23 @@ describe('Tree', () => {
     expect(wrapper.render()).toMatchSnapshot();
   });
 
+  it('switcherIcon in Tree could be render prop function', () => {
+    const wrapper = mount(
+      <Tree
+        switcherIcon={expanded =>
+          expanded ? <span className="open" /> : <span className="close" />
+        }
+        defaultExpandAll
+      >
+        <TreeNode icon="icon">
+          <TreeNode id="node1" title="node1" icon="icon" key="0-0-2" />
+          <TreeNode id="node2" title="node2" key="0-0-3" />
+        </TreeNode>
+      </Tree>,
+    );
+    expect(wrapper.find('.open').length).toBe(1);
+  });
+
   // https://github.com/ant-design/ant-design/issues/23261
   it('showLine is object type should render correctly', () => {
     const wrapper = mount(

components/tree/index.en-US.md

@@ -44,7 +44,7 @@ Almost anything can be represented in a tree structure. Examples include directo
 | selectedKeys | (Controlled) Specifies the keys of the selected treeNodes | string\[] | - |  |
 | showIcon | Shows the icon before a TreeNode's title. There is no default style; you must set a custom style for it if set to true | boolean | false |  |
 | showLine | Shows a connecting line | boolean \| {showLeafIcon: boolean} | false |  |
-| switcherIcon | Customize collapse/expand icon of tree node | ReactNode | - |  |
+| switcherIcon | Customize collapse/expand icon of tree node | ReactNode \|  ((expanded: boolean) => React.ReactNode) | - |  |
 | titleRender | Customize tree node title render | (nodeData) => ReactNode | - | 4.5.0 |
 | treeData | The treeNodes data Array, if set it then you need not to construct children TreeNode. (key should be unique across the whole array) | array<{ key, title, children, \[disabled, selectable] }> | - |  |
 | virtual | Disable virtual scroll when set to false | boolean | true | 4.1.0 |

components/tree/index.zh-CN.md

@@ -45,7 +45,7 @@ cover: https://gw.alipayobjects.com/zos/alicdn/Xh-oWqg9k/Tree.svg
 | selectedKeys | （受控）设置选中的树节点 | string\[] | - |  |
 | showIcon | 是否展示 TreeNode title 前的图标，没有默认样式，如设置为 true，需要自行定义图标相关样式 | boolean | false |  |
 | showLine | 是否展示连接线 | boolean \| {showLeafIcon: boolean} | false |  |
-| switcherIcon | 自定义树节点的展开/折叠图标 | ReactNode | - |  |
+| switcherIcon | 自定义树节点的展开/折叠图标 | ReactNode \|  ((expanded: boolean) => React.ReactNode) | - |  |
 | titleRender | 自定义渲染节点 | (nodeData) => ReactNode | - | 4.5.0 |
 | treeData | treeNodes 数据，如果设置则不需要手动构造 TreeNode 节点（key 在整个树范围内唯一） | array<{key, title, children, \[disabled, selectable]}> | - |  |
 | virtual | 设置 false 时关闭虚拟滚动 | boolean | true | 4.1.0 |

components/tree/utils/iconUtil.tsx

@@ -5,12 +5,12 @@ import FileOutlined from '@ant-design/icons/FileOutlined';
 import MinusSquareOutlined from '@ant-design/icons/MinusSquareOutlined';
 import PlusSquareOutlined from '@ant-design/icons/PlusSquareOutlined';
 import CaretDownFilled from '@ant-design/icons/CaretDownFilled';
-import { AntTreeNodeProps } from '../Tree';
+import { AntTreeNodeProps, SwitcherIcon } from '../Tree';
 import { isValidElement, cloneElement } from '../../_util/reactNode';
 
 export default function renderSwitcherIcon(
   prefixCls: string,
-  switcherIcon: React.ReactNode,
+  switcherIcon: SwitcherIcon | undefined,
   showLine: boolean | { showLeafIcon: boolean } | undefined,
   { isLeaf, expanded, loading }: AntTreeNodeProps,
 ) {
@@ -31,14 +31,20 @@ export default function renderSwitcherIcon(
     return null;
   }
   const switcherCls = `${prefixCls}-switcher-icon`;
-  if (isValidElement(switcherIcon)) {
-    return cloneElement(switcherIcon, {
-      className: classNames(switcherIcon.props.className || '', switcherCls),
+
+  let switcher = switcherIcon;
+  if (typeof switcher === 'function') {
+    switcher = switcher(Boolean(expanded));
+  }
+
+  if (isValidElement(switcher)) {
+    return cloneElement(switcher, {
+      className: classNames(switcher.props.className || '', switcherCls),
     });
   }
 
-  if (switcherIcon) {
-    return switcherIcon;
+  if (switcher) {
+    return switcher;
   }
 
   if (showLine) {