feat(tree-select): 暴露 treeRef 的方法 (#1698)

fix #1695
Tencent · Nov 10, 2022 · fcb046c · fcb046c
1 parent 164cbe1
commit fcb046c
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/src/tree-select/TreeSelect.tsx b/src/tree-select/TreeSelect.tsx
@@ -1,4 +1,4 @@
-import React, { useCallback, useMemo, useRef, forwardRef, ElementRef, useEffect } from 'react';
+import React, { useCallback, useMemo, useRef, forwardRef, ElementRef, useEffect, useImperativeHandle } from 'react';
 
 import classNames from 'classnames';
 import type { TdTreeSelectProps, TreeSelectValue } from './type';
@@ -28,7 +28,7 @@ export interface NodeOptions {
 const useMergeFn = <T extends any[]>(...fns: Array<(...args: T) => void>) =>
   usePersistFn((...args: T) => fns.forEach((fn) => fn?.(...args)));
 
-const TreeSelect = forwardRef((props: TreeSelectProps, ref: React.Ref<HTMLDivElement>) => {
+const TreeSelect = forwardRef((props: TreeSelectProps, ref) => {
   /* ---------------------------------config---------------------------------------- */
 
   // 国际化文本初始化
@@ -68,9 +68,15 @@ const TreeSelect = forwardRef((props: TreeSelectProps, ref: React.Ref<HTMLDivEle
   const [filterInput, setFilterInput] = useControlled(props, 'inputValue', onInputChange);
 
   const treeRef = useRef<ElementRef<typeof Tree>>();
+  const selectInputRef = useRef();
 
   const { normalizeValue, formatValue, getNodeItem } = useTreeSelectUtils(props, treeRef);
 
+  useImperativeHandle(ref, () => ({
+    ...(selectInputRef.current || {}),
+    ...(treeRef.current || {}),
+  }));
+
   /* ---------------------------------computed value---------------------------------------- */
 
   const filterable = rawFilterable || !!props.filter;
@@ -273,7 +279,7 @@ const TreeSelect = forwardRef((props: TreeSelectProps, ref: React.Ref<HTMLDivEle
       tips={props.tips}
       {...props.selectInputProps}
       {...selectInputProps}
-      ref={ref}
+      ref={selectInputRef}
       className={classNames(
         `${classPrefix}-tree-select`,
         {