Merge pull request #2237 from DouyinFE/feat/tree-custom-icon

feat: [Tree] icon API support function

content/navigation/tree/index-en-US.md

@@ -2288,7 +2288,7 @@ import { IconFixedStroked, IconSectionStroked, IconAbsoluteStroked, IconInnerSec
 | keyMaps | Customize the key, label, and value fields in the node | object |  - | 2.47.0 |
 | filterTreeNode      | Toggle whether searchable or pass in a function to customize search behavior, data parameter provided since v2.28.0 | boolean \| ((inputValue: string, treeNodeString: string, data?: TreeNodeData) => boolean)  | false   | - |
 | hideDraggingNode | Toggle whether to hide dragImg of dragging node | boolean | false | 1.8.0 | 
-| icon       | Icon | ReactNode         | -       | - |
+| icon       | Icon | ReactNode \|(props: TreeNodeProps) => ReactNode         | -       | - |
 | labelEllipsis | Toggle whether to ellipsis label when overflow. Set to false iff there are other requirements | boolean | false\|true(virtualized) | 1.8.0 | 
 | leafOnly | Toggle whether to display tags for leaf nodes only and for onChange callback params in multiple mode | boolean | false | 1.18.0 |
 | loadData | Load data asynchronously and the return value should be a promise | (treeNode?: TreeNodeData) => Promise< void > |-| 1.0.0|

content/navigation/tree/index.md

@@ -2303,7 +2303,7 @@ import { IconFixedStroked, IconSectionStroked, IconAbsoluteStroked, IconInnerSec
 | keyMaps | 自定义节点中 key、label、value 的字段 | object |  - | 2.47.0 |
 | filterTreeNode | 是否根据输入项进行筛选，默认用 `treeNodeFilterProp` 的值作为要筛选的 `TreeNodeData` 的属性值,  data 参数自 v2.28.0 开始提供 | boolean \| ((inputValue: string, treeNodeString: string, data?: TreeNodeData) => boolean) | false | - |
 | hideDraggingNode | 是否隐藏正在拖拽的节点的 dragImg | boolean | false | 1.8.0 | 
-| icon | 自定义图标 | ReactNode | - | - |
+| icon | 自定义图标 | ReactNode \| (props: TreeNodeProps)=>ReactNode | - | - |
 | labelEllipsis | 是否开启label的超出省略，默认虚拟化状态开启，如果有其他省略需求可以设置关闭 | boolean | false\|true(virtualized) | 1.8.0 |
 | leafOnly | 多选模式下是否开启 onChange 回调入参及展示标签只有叶子节点 | boolean | false | 1.18.0 |  
 | loadData | 异步加载数据，需要返回一个Promise | (treeNode?: TreeNodeData) => Promise< void > |- |  1.0.0|

packages/semi-ui/tree/_story/tree.stories.jsx

@@ -3,7 +3,7 @@ import { difference, isEqual } from 'lodash';
 import { IconEdit, IconMapPin, IconMore } from '@douyinfe/semi-icons';
 import Tree from '../index';
 import AutoSizer from '../autoSizer';
-import { Button, ButtonGroup, Input, Popover, Toast, Space, Select, Switch, Typography } from '../../index';
+import { Button, ButtonGroup, Input, Popover, Toast, Space, Select, Switch, Typography, Tag } from '../../index';
 import BigTree from './BigData';
 import testData from './data';
 import copy from 'fast-copy';
@@ -3085,3 +3085,35 @@ export const ShowLine = () => {
 ShowLine.story = {
   name: 'show line',
 };
+
+export const CustomRenderIcon = () => {
+  const iconFunc = useCallback((props) => {
+    const { data } = props;
+    return data.suffix ? <Tag>{data.suffix}</Tag> : null
+  })
+
+  const treeDataWithSuffix = [
+    {
+      label: 'package',
+      key: '0',
+      children: [
+        {
+          label: 'semi.py',
+          suffix: 'Py',
+          key: '1'
+        },
+        {
+          label: 'semi.js',
+          suffix: 'Js',
+          key: '2'
+        }
+      ]
+    }, 
+  ]
+
+  return (<Tree
+      defaultExpandAll
+      treeData={treeDataWithSuffix}
+      icon={iconFunc}
+  />);
+}

packages/semi-ui/tree/interface.ts

@@ -67,7 +67,7 @@ export interface TreeProps extends BasicTreeProps {
     style?: React.CSSProperties;
     treeData?: TreeNodeData[];
     value?: Value;
-    icon?: ReactNode;
+    icon?: ReactNode | ((props: TreeNodeProps) => ReactNode);
     keyMaps?: KeyMapProps;
     loadData?: (treeNode?: TreeNodeData) => Promise<void>;
     onChange?: (value?: Value) => void;

packages/semi-ui/tree/treeContext.tsx

@@ -10,7 +10,7 @@ import {
 
 export interface TreeContextValue {
     treeDisabled?: boolean;
-    treeIcon?: ReactNode;
+    treeIcon?: ReactNode | ((props: TreeNodeProps) => ReactNode);
     motion?: boolean;
     motionKeys?: Set<string>;
     motionType?: string;

packages/semi-ui/tree/treeNode.tsx

@@ -263,22 +263,22 @@ export default class TreeNode extends PureComponent<TreeNodeProps, TreeNodeState
             directory,
             treeIcon
         } = this.context;
-        const { expanded, icon } = this.props;
-        const hasChild = !this.isLeaf();
-        const hasIcon = icon || treeIcon;
-        let itemIcon;
-        if (hasIcon || directory) {
-            if (hasIcon) {
-                itemIcon = icon || treeIcon;
+        const { expanded, icon, data } = this.props;
+        if (icon) {
+            return icon;
+        }
+        if (treeIcon) {
+            return typeof treeIcon === 'function' ? treeIcon(this.props) : treeIcon;
+        }
+        if (directory) {
+            const hasChild = !this.isLeaf();
+            if (!hasChild) {
+                return <IconFile className={`${prefixcls}-item-icon`} />;
             } else {
-                if (!hasChild) {
-                    itemIcon = <IconFile className={`${prefixcls}-item-icon`} />;
-                } else {
-                    itemIcon = expanded ? <IconFolderOpen className={`${prefixcls}-item-icon`} /> : <IconFolder className={`${prefixcls}-item-icon`} />;
-                }
+                return expanded ? <IconFolderOpen className={`${prefixcls}-item-icon`} /> : <IconFolder className={`${prefixcls}-item-icon`} />;
             }
         }
-        return itemIcon;
+        return null;
     }
 
     renderEmptyNode() {