refactor(tree-select): rewrite with hook (#27593)

components/tree-select/__tests__/index.test.js

@@ -6,7 +6,7 @@ import mountTest from '../../../tests/shared/mountTest';
 import rtlTest from '../../../tests/shared/rtlTest';
 
 describe('TreeSelect', () => {
-  focusTest(TreeSelect);
+  focusTest(TreeSelect, { refFocus: true });
   mountTest(TreeSelect);
   rtlTest(TreeSelect);
 

components/tree-select/index.tsx

@@ -8,7 +8,8 @@ import RcTreeSelect, {
 } from 'rc-tree-select';
 import classNames from 'classnames';
 import omit from 'omit.js';
-import { ConfigConsumer, ConfigConsumerProps } from '../config-provider';
+import { DefaultValueType } from 'rc-tree-select/lib/interface';
+import { ConfigContext } from '../config-provider';
 import devWarning from '../_util/devWarning';
 import { AntTreeNodeProps } from '../tree';
 import getIcons from '../select/utils/iconUtil';
@@ -35,167 +36,147 @@ export interface TreeSelectProps<T>
   bordered?: boolean;
 }
 
-class TreeSelect<T> extends React.Component<TreeSelectProps<T>, {}> {
-  static TreeNode = TreeNode;
+export interface RefTreeSelectProps {
+  focus: () => void;
+  blur: () => void;
+}
 
-  static SHOW_ALL: typeof SHOW_ALL = SHOW_ALL;
-
-  static SHOW_PARENT: typeof SHOW_PARENT = SHOW_PARENT;
-
-  static SHOW_CHILD: typeof SHOW_CHILD = SHOW_CHILD;
-
-  static defaultProps = {
-    transitionName: 'slide-up',
-    choiceTransitionName: '',
-    bordered: true,
-  };
-
-  selectRef = React.createRef<RcTreeSelect>();
-
-  constructor(props: TreeSelectProps<T>) {
-    super(props);
-
-    devWarning(
-      props.multiple !== false || !props.treeCheckable,
-      'TreeSelect',
-      '`multiple` will alway be `true` when `treeCheckable` is true',
-    );
-  }
-
-  focus() {
-    if (this.selectRef.current) {
-      this.selectRef.current.focus();
-    }
-  }
-
-  blur() {
-    if (this.selectRef.current) {
-      this.selectRef.current.blur();
-    }
-  }
-
-  renderTreeSelect = ({
+const InternalTreeSelect = <T extends DefaultValueType>(
+  {
+    prefixCls: customizePrefixCls,
+    size: customizeSize,
+    bordered = true,
+    className,
+    treeCheckable,
+    multiple,
+    listHeight = 256,
+    listItemHeight = 26,
+    notFoundContent,
+    switcherIcon,
+    treeLine,
+    getPopupContainer,
+    dropdownClassName,
+    treeIcon = false,
+    transitionName = 'slide-up',
+    choiceTransitionName = '',
+    ...props
+  }: TreeSelectProps<T>,
+  ref: React.Ref<RefTreeSelectProps>,
+) => {
+  const {
     getPopupContainer: getContextPopupContainer,
     getPrefixCls,
     renderEmpty,
     direction,
     virtual,
     dropdownMatchSelectWidth,
-  }: ConfigConsumerProps) => {
-    const {
-      prefixCls: customizePrefixCls,
-      size: customizeSize,
-      className,
-      treeCheckable,
-      multiple,
-      listHeight = 256,
-      listItemHeight = 26,
-      notFoundContent,
-      switcherIcon,
-      treeLine,
-      getPopupContainer,
-      dropdownClassName,
-      bordered,
-      treeIcon = false,
-    } = this.props;
+  } = React.useContext(ConfigContext);
+  const size = React.useContext(SizeContext);
 
-    const prefixCls = getPrefixCls('select', customizePrefixCls);
-    const treePrefixCls = getPrefixCls('select-tree', customizePrefixCls);
-    const treeSelectPrefixCls = getPrefixCls('tree-select', customizePrefixCls);
+  devWarning(
+    multiple !== false || !treeCheckable,
+    'TreeSelect',
+    '`multiple` will alway be `true` when `treeCheckable` is true',
+  );
 
-    const mergedDropdownClassName = classNames(
-      dropdownClassName,
-      `${treeSelectPrefixCls}-dropdown`,
-      {
-        [`${treeSelectPrefixCls}-dropdown-rtl`]: direction === 'rtl',
-      },
-    );
+  const prefixCls = getPrefixCls('select', customizePrefixCls);
+  const treePrefixCls = getPrefixCls('select-tree', customizePrefixCls);
+  const treeSelectPrefixCls = getPrefixCls('tree-select', customizePrefixCls);
 
-    const isMultiple = !!(treeCheckable || multiple);
+  const mergedDropdownClassName = classNames(dropdownClassName, `${treeSelectPrefixCls}-dropdown`, {
+    [`${treeSelectPrefixCls}-dropdown-rtl`]: direction === 'rtl',
+  });
 
-    // ===================== Icons =====================
-    const { suffixIcon, itemIcon, removeIcon, clearIcon } = getIcons({
-      ...this.props,
-      multiple: isMultiple,
-      prefixCls,
-    });
+  const isMultiple = !!(treeCheckable || multiple);
 
-    // ===================== Empty =====================
-    let mergedNotFound: React.ReactNode;
-    if (notFoundContent !== undefined) {
-      mergedNotFound = notFoundContent;
-    } else {
-      mergedNotFound = renderEmpty('Select');
-    }
+  // ===================== Icons =====================
+  const { suffixIcon, itemIcon, removeIcon, clearIcon } = getIcons({
+    ...props,
+    multiple: isMultiple,
+    prefixCls,
+  });
 
-    // ==================== Render =====================
-    const selectProps = omit(this.props, [
-      'prefixCls',
-      'suffixIcon',
-      'itemIcon',
-      'removeIcon',
-      'clearIcon',
-      'switcherIcon',
-      'size',
-      'bordered',
-    ]);
-
-    return (
-      <SizeContext.Consumer>
-        {size => {
-          const mergedSize = customizeSize || size;
-          const mergedClassName = classNames(
-            !customizePrefixCls && treeSelectPrefixCls,
-            {
-              [`${prefixCls}-lg`]: mergedSize === 'large',
-              [`${prefixCls}-sm`]: mergedSize === 'small',
-              [`${prefixCls}-rtl`]: direction === 'rtl',
-              [`${prefixCls}-borderless`]: !bordered,
-            },
-            className,
-          );
-
-          return (
-            <RcTreeSelect
-              virtual={virtual}
-              dropdownMatchSelectWidth={dropdownMatchSelectWidth}
-              {...selectProps}
-              ref={this.selectRef}
-              prefixCls={prefixCls}
-              className={mergedClassName}
-              listHeight={listHeight}
-              listItemHeight={listItemHeight}
-              treeCheckable={
-                treeCheckable ? (
-                  <span className={`${prefixCls}-tree-checkbox-inner`} />
-                ) : (
-                  treeCheckable
-                )
-              }
-              inputIcon={suffixIcon}
-              menuItemSelectedIcon={itemIcon}
-              removeIcon={removeIcon}
-              clearIcon={clearIcon}
-              switcherIcon={(nodeProps: AntTreeNodeProps) =>
-                renderSwitcherIcon(treePrefixCls, switcherIcon, treeLine, nodeProps)
-              }
-              showTreeIcon={treeIcon}
-              notFoundContent={mergedNotFound}
-              getPopupContainer={getPopupContainer || getContextPopupContainer}
-              treeMotion={null}
-              dropdownClassName={mergedDropdownClassName}
-            />
-          );
-        }}
-      </SizeContext.Consumer>
-    );
-  };
-
-  render() {
-    return <ConfigConsumer>{this.renderTreeSelect}</ConfigConsumer>;
+  // ===================== Empty =====================
+  let mergedNotFound: React.ReactNode;
+  if (notFoundContent !== undefined) {
+    mergedNotFound = notFoundContent;
+  } else {
+    mergedNotFound = renderEmpty('Select');
   }
+
+  // ==================== Render =====================
+  const selectProps = omit(props, [
+    'suffixIcon',
+    'itemIcon',
+    'removeIcon',
+    'clearIcon',
+    'switcherIcon',
+  ]);
+
+  const mergedSize = customizeSize || size;
+  const mergedClassName = classNames(
+    !customizePrefixCls && treeSelectPrefixCls,
+    {
+      [`${prefixCls}-lg`]: mergedSize === 'large',
+      [`${prefixCls}-sm`]: mergedSize === 'small',
+      [`${prefixCls}-rtl`]: direction === 'rtl',
+      [`${prefixCls}-borderless`]: !bordered,
+    },
+    className,
+  );
+
+  return (
+    <RcTreeSelect
+      virtual={virtual}
+      dropdownMatchSelectWidth={dropdownMatchSelectWidth}
+      {...selectProps}
+      ref={ref}
+      prefixCls={prefixCls}
+      className={mergedClassName}
+      listHeight={listHeight}
+      listItemHeight={listItemHeight}
+      treeCheckable={
+        treeCheckable ? <span className={`${prefixCls}-tree-checkbox-inner`} /> : treeCheckable
+      }
+      inputIcon={suffixIcon}
+      menuItemSelectedIcon={itemIcon}
+      multiple={multiple}
+      removeIcon={removeIcon}
+      clearIcon={clearIcon}
+      switcherIcon={(nodeProps: AntTreeNodeProps) =>
+        renderSwitcherIcon(treePrefixCls, switcherIcon, treeLine, nodeProps)
+      }
+      showTreeIcon={treeIcon}
+      notFoundContent={mergedNotFound}
+      getPopupContainer={getPopupContainer || getContextPopupContainer}
+      treeMotion={null}
+      dropdownClassName={mergedDropdownClassName}
+      choiceTransitionName={choiceTransitionName}
+      transitionName={transitionName}
+    />
+  );
+};
+
+const TreeSelectRef = React.forwardRef(InternalTreeSelect) as <T extends DefaultValueType>(
+  props: TreeSelectProps<T> & { ref?: React.Ref<RefTreeSelectProps> },
+) => React.ReactElement;
+
+type InternalTreeSelectType = typeof TreeSelectRef;
+
+interface TreeSelectInterface extends InternalTreeSelectType {
+  TreeNode: typeof TreeNode;
+  SHOW_ALL: typeof SHOW_ALL;
+  SHOW_PARENT: typeof SHOW_PARENT;
+  SHOW_CHILD: typeof SHOW_CHILD;
 }
 
+const TreeSelect = TreeSelectRef as TreeSelectInterface;
+
+TreeSelect.TreeNode = TreeNode;
+TreeSelect.SHOW_ALL = SHOW_ALL;
+TreeSelect.SHOW_PARENT = SHOW_PARENT;
+TreeSelect.SHOW_CHILD = SHOW_CHILD;
+
 export { TreeNode };
 
 export default TreeSelect;