From 5f5be87a4d90987b42b560df17b338a9c6493c73 Mon Sep 17 00:00:00 2001
From: huangkairan <56213366+huangkairan@users.noreply.github.com>
Date: Mon, 4 Dec 2023 15:30:33 +0800
Subject: [PATCH] feat: support titleRender (#447)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* ✨ feat: support titleRender

* 🐳 chore: update test case

* 🐳 chore: update prop name

* 🐳 chore: update `label` logic

* 🌈 style: format
---
 examples/basic.tsx         | 16 ++++++++++++----
 src/OptionList.tsx         |  2 ++
 src/TreeSelect.tsx         | 13 ++++++++++++-
 src/TreeSelectContext.ts   |  1 +
 tests/Select.props.spec.js | 22 +++++++++++++++++++++-
 5 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/examples/basic.tsx b/examples/basic.tsx
index bc96622e..a0d97993 100644
--- a/examples/basic.tsx
+++ b/examples/basic.tsx
@@ -1,8 +1,8 @@
-import '../assets/index.less';
-import React from 'react';
-import 'rc-dialog/assets/index.css';
 import Dialog from 'rc-dialog';
-import TreeSelect, { TreeNode, SHOW_PARENT } from '../src';
+import 'rc-dialog/assets/index.css';
+import React from 'react';
+import '../assets/index.less';
+import TreeSelect, { SHOW_PARENT, TreeNode } from '../src';
 import { gData } from './utils/dataUtil';
 
 function isLeaf(value) {
@@ -381,6 +381,14 @@ class Demo extends React.Component {
           </TreeNode>
           <TreeNode value="same value3" title="same title" key="0-3" />
         </TreeSelect>
+
+        <h2>title render</h2>
+        <TreeSelect<{ label: string }>
+          open
+          style={{ width: 300 }}
+          treeData={gData}
+          treeTitleRender={node => node.label + 'ok'}
+        />
       </div>
     );
   }
diff --git a/src/OptionList.tsx b/src/OptionList.tsx
index d96ac542..6c83c982 100644
--- a/src/OptionList.tsx
+++ b/src/OptionList.tsx
@@ -43,6 +43,7 @@ const OptionList: React.RefForwardingComponent<ReviseRefOptionListProps> = (_, r
     onSelect,
     dropdownMatchSelectWidth,
     treeExpandAction,
+    treeTitleRender,
   } = React.useContext(TreeSelectContext);
 
   const {
@@ -241,6 +242,7 @@ const OptionList: React.RefForwardingComponent<ReviseRefOptionListProps> = (_, r
         checkedKeys={mergedCheckedKeys}
         selectedKeys={!checkable ? checkedKeys : []}
         defaultExpandAll={treeDefaultExpandAll}
+        titleRender={treeTitleRender}
         {...treeProps}
         // Proxy event out
         onActiveChange={setActiveKey}
diff --git a/src/TreeSelect.tsx b/src/TreeSelect.tsx
index 1fb27b85..051cdddf 100644
--- a/src/TreeSelect.tsx
+++ b/src/TreeSelect.tsx
@@ -161,6 +161,7 @@ export interface TreeSelectProps<
   listItemHeight?: number;
   listItemScrollOffset?: number;
   onDropdownVisibleChange?: (open: boolean) => void;
+  treeTitleRender?: (node: ValueType) => React.ReactNode;
 
   // >>> Tree
   treeLine?: boolean;
@@ -237,6 +238,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
     showTreeIcon,
     switcherIcon,
     treeMotion,
+    treeTitleRender,
 
     ...restProps
   } = props;
@@ -438,9 +440,16 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
     // Back fill with origin label
     const labeledValues = values.map(val => {
       const targetItem = rawLabeledValues.find(item => item.value === val);
+      let label;
+      // Ensure that when labelInValue is true, if label is undefined, it remains undefined.
+      if (labelInValue && targetItem.label !== undefined) {
+        label = targetItem.label;
+      } else if (!labelInValue && treeTitleRender) {
+        label = treeTitleRender(targetItem);
+      }
       return {
         value: val,
-        label: targetItem?.label,
+        label,
       };
     });
 
@@ -671,6 +680,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       fieldNames: mergedFieldNames,
       onSelect: onOptionSelect,
       treeExpandAction,
+      treeTitleRender,
     }),
     [
       virtual,
@@ -682,6 +692,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       mergedFieldNames,
       onOptionSelect,
       treeExpandAction,
+      treeTitleRender,
     ],
   );
 
diff --git a/src/TreeSelectContext.ts b/src/TreeSelectContext.ts
index bcdb223a..344f71de 100644
--- a/src/TreeSelectContext.ts
+++ b/src/TreeSelectContext.ts
@@ -12,6 +12,7 @@ export interface TreeSelectContextProps {
   fieldNames: InternalFieldName;
   onSelect: OnInternalSelect;
   treeExpandAction?: ExpandAction;
+  treeTitleRender?: (node: any) => React.ReactNode;
 }
 
 const TreeSelectContext = React.createContext<TreeSelectContextProps>(null as any);
diff --git a/tests/Select.props.spec.js b/tests/Select.props.spec.js
index cb5c38dc..81ef4fe7 100644
--- a/tests/Select.props.spec.js
+++ b/tests/Select.props.spec.js
@@ -1,7 +1,7 @@
 /* eslint-disable no-undef, react/no-multi-comp, no-console */
-import React from 'react';
 import { mount } from 'enzyme';
 import Tree, { TreeNode } from 'rc-tree';
+import React from 'react';
 import TreeSelect, { SHOW_ALL, SHOW_CHILD, SHOW_PARENT, TreeNode as SelectNode } from '../src';
 
 // Promisify timeout to let jest catch works
@@ -612,5 +612,25 @@ describe('TreeSelect.props', () => {
         });
       });
     });
+
+    describe('title render', () => {
+      const treeData = [
+        { label: 'Label 0-0', value: 'Value 0-0', key: 'key 0-0' },
+        { label: 'Label 0-1', value: 'Value 0-1', key: 'key 0-1' },
+        { label: 'Label 1-0', value: 'Value 1-0', key: 'key 1-0' },
+      ];
+      it('basic', () => {
+        const wrapper = mount(
+          <div>
+            <TreeSelect
+              defaultValue={'Value 0-0'}
+              treeTitleRender={node => node.value}
+              treeData={treeData}
+            />
+          </div>,
+        );
+        expect(wrapper.getSelection(0).text()).toBe('Value 0-0');
+      });
+    });
   });
 });