From 11f81f3032a5b7e59a5fe7167317c16d34174c54 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=8C=E8=B4=A7=E6=9C=BA=E5=99=A8=E4=BA=BA?=
 <smith3816@gmail.com>
Date: Tue, 22 Jun 2021 15:52:11 +0800
Subject: [PATCH] feat: Add getRawInputElement (#639)

* feat: Add raw customize

* feat: Support dropdown

* test: Add test case
---
 examples/custom-selector.tsx |  17 +++++
 src/SelectTrigger.tsx        |   7 +-
 src/generate.tsx             | 124 ++++++++++++++++++++++-------------
 tests/Select.test.tsx        |  18 +++++
 typings/index.d.ts           |   2 -
 5 files changed, 118 insertions(+), 50 deletions(-)
 create mode 100644 examples/custom-selector.tsx

diff --git a/examples/custom-selector.tsx b/examples/custom-selector.tsx
new file mode 100644
index 000000000..381cc2e22
--- /dev/null
+++ b/examples/custom-selector.tsx
@@ -0,0 +1,17 @@
+/* eslint-disable no-console */
+import React from 'react';
+import Select from '../src';
+import '../assets/index.less';
+
+export default () => {
+  return (
+    <Select
+      getRawInputElement={() => <span>Content</span>}
+      mode="multiple"
+      options={[{ value: 'light' }, { value: 'bamboo' }]}
+      allowClear
+      placeholder="2333"
+    />
+  );
+};
+/* eslint-enable */
diff --git a/src/SelectTrigger.tsx b/src/SelectTrigger.tsx
index 821139528..cad9d1bfc 100644
--- a/src/SelectTrigger.tsx
+++ b/src/SelectTrigger.tsx
@@ -67,6 +67,7 @@ export interface SelectTriggerProps {
   empty: boolean;
 
   getTriggerDOMNode: () => HTMLElement;
+  onPopupVisibleChange?: (visible: boolean) => void;
 }
 
 const SelectTrigger: React.RefForwardingComponent<RefTriggerProps, SelectTriggerProps> = (
@@ -91,6 +92,7 @@ const SelectTrigger: React.RefForwardingComponent<RefTriggerProps, SelectTrigger
     getPopupContainer,
     empty,
     getTriggerDOMNode,
+    onPopupVisibleChange,
     ...restProps
   } = props;
 
@@ -129,8 +131,8 @@ const SelectTrigger: React.RefForwardingComponent<RefTriggerProps, SelectTrigger
   return (
     <Trigger
       {...restProps}
-      showAction={[]}
-      hideAction={[]}
+      showAction={onPopupVisibleChange ? ['click'] : []}
+      hideAction={onPopupVisibleChange ? ['click'] : []}
       popupPlacement={direction === 'rtl' ? 'bottomRight' : 'bottomLeft'}
       builtinPlacements={builtInPlacements}
       prefixCls={dropdownPrefixCls}
@@ -144,6 +146,7 @@ const SelectTrigger: React.RefForwardingComponent<RefTriggerProps, SelectTrigger
       })}
       popupStyle={popupStyle}
       getTriggerDOMNode={getTriggerDOMNode}
+      onPopupVisibleChange={onPopupVisibleChange}
     >
       {children}
     </Trigger>
diff --git a/src/generate.tsx b/src/generate.tsx
index 0721e8fab..18ecaa98e 100644
--- a/src/generate.tsx
+++ b/src/generate.tsx
@@ -11,6 +11,7 @@ import * as React from 'react';
 import { useState, useRef, useEffect, useMemo } from 'react';
 import KeyCode from 'rc-util/lib/KeyCode';
 import isMobile from 'rc-util/lib/isMobile';
+import { composeRef } from 'rc-util/lib/ref';
 import classNames from 'classnames';
 import useMergedState from 'rc-util/lib/hooks/useMergedState';
 import type { ScrollTo } from 'rc-virtual-list/lib/List';
@@ -133,6 +134,8 @@ export interface SelectProps<OptionsType extends object[], ValueType> extends Re
   backfill?: boolean;
   /** @private Internal usage. Do not use in your production. */
   getInputElement?: () => JSX.Element;
+  /** @private Internal usage. Do not use in your production. */
+  getRawInputElement?: () => JSX.Element;
   optionLabelProp?: string;
   maxTagTextLength?: number;
   maxTagCount?: number | 'responsive';
@@ -289,6 +292,7 @@ export default function generateSelector<
       backfill,
       tabIndex,
       getInputElement,
+      getRawInputElement,
       getPopupContainer,
 
       // Dropdown
@@ -635,9 +639,13 @@ export default function generateSelector<
     };
 
     // ============================= Input ==============================
-    // Only works in `combobox` or `rc-cascader`
+    // Only works in `combobox`
     const customizeInputElement: React.ReactElement =
-      (typeof getInputElement === 'function' && getInputElement()) || null;
+      (mode === 'combobox' && typeof getInputElement === 'function' && getInputElement()) || null;
+
+    // Used for customize replacement for `rc-cascader`
+    const customizeRawInputElement: React.ReactElement =
+      typeof getRawInputElement === 'function' && getRawInputElement();
 
     // ============================== Open ==============================
     const [innerOpen, setInnerOpen] = useMergedState<boolean>(undefined, {
@@ -666,6 +674,14 @@ export default function generateSelector<
       }
     };
 
+    // Used for raw custom input trigger
+    let onTriggerVisibleChange: null | ((newOpen: boolean) => void);
+    if (customizeRawInputElement) {
+      onTriggerVisibleChange = (newOpen: boolean) => {
+        onToggleOpen(newOpen);
+      };
+    }
+
     useSelectTriggerControl(
       [containerRef.current, triggerRef.current && triggerRef.current.getPopupElement()],
       triggerOpen,
@@ -931,8 +947,8 @@ export default function generateSelector<
 
     useLayoutEffect(() => {
       if (triggerOpen) {
-        const newWidth = Math.ceil(containerRef.current.offsetWidth);
-        if (containerWidth !== newWidth) {
+        const newWidth = Math.ceil(containerRef.current?.offsetWidth);
+        if (containerWidth !== newWidth && !Number.isNaN(newWidth)) {
           setContainerWidth(newWidth);
         }
       }
@@ -1034,6 +1050,63 @@ export default function generateSelector<
       [`${prefixCls}-show-search`]: mergedShowSearch,
     });
 
+    const selectorNode = (
+      <SelectTrigger
+        ref={triggerRef}
+        disabled={disabled}
+        prefixCls={prefixCls}
+        visible={triggerOpen}
+        popupElement={popupNode}
+        containerWidth={containerWidth}
+        animation={animation}
+        transitionName={transitionName}
+        dropdownStyle={dropdownStyle}
+        dropdownClassName={dropdownClassName}
+        direction={direction}
+        dropdownMatchSelectWidth={dropdownMatchSelectWidth}
+        dropdownRender={dropdownRender}
+        dropdownAlign={dropdownAlign}
+        getPopupContainer={getPopupContainer}
+        empty={!mergedOptions.length}
+        getTriggerDOMNode={() => selectorDomRef.current}
+        onPopupVisibleChange={onTriggerVisibleChange}
+      >
+        {customizeRawInputElement ? (
+          React.cloneElement(customizeRawInputElement, {
+            ref: composeRef(selectorDomRef, customizeRawInputElement.props.ref),
+          })
+        ) : (
+          <Selector
+            {...props}
+            domRef={selectorDomRef}
+            prefixCls={prefixCls}
+            inputElement={customizeInputElement}
+            ref={selectorRef}
+            id={mergedId}
+            showSearch={mergedShowSearch}
+            mode={mode}
+            accessibilityIndex={accessibilityIndex}
+            multiple={isMultiple}
+            tagRender={tagRender}
+            values={displayValues}
+            open={mergedOpen}
+            onToggleOpen={onToggleOpen}
+            searchValue={mergedSearchValue}
+            activeValue={activeValue}
+            onSearch={triggerSearch}
+            onSearchSubmit={onSearchSubmit}
+            onSelect={onInternalSelectionSelect}
+            tokenWithEnter={tokenWithEnter}
+          />
+        )}
+      </SelectTrigger>
+    );
+
+    // Render raw
+    if (customizeRawInputElement) {
+      return selectorNode;
+    }
+
     return (
       <div
         className={mergedClassName}
@@ -1060,48 +1133,7 @@ export default function generateSelector<
             {`${mergedRawValue.join(', ')}`}
           </span>
         )}
-        <SelectTrigger
-          ref={triggerRef}
-          disabled={disabled}
-          prefixCls={prefixCls}
-          visible={triggerOpen}
-          popupElement={popupNode}
-          containerWidth={containerWidth}
-          animation={animation}
-          transitionName={transitionName}
-          dropdownStyle={dropdownStyle}
-          dropdownClassName={dropdownClassName}
-          direction={direction}
-          dropdownMatchSelectWidth={dropdownMatchSelectWidth}
-          dropdownRender={dropdownRender}
-          dropdownAlign={dropdownAlign}
-          getPopupContainer={getPopupContainer}
-          empty={!mergedOptions.length}
-          getTriggerDOMNode={() => selectorDomRef.current}
-        >
-          <Selector
-            {...props}
-            domRef={selectorDomRef}
-            prefixCls={prefixCls}
-            inputElement={customizeInputElement}
-            ref={selectorRef}
-            id={mergedId}
-            showSearch={mergedShowSearch}
-            mode={mode}
-            accessibilityIndex={accessibilityIndex}
-            multiple={isMultiple}
-            tagRender={tagRender}
-            values={displayValues}
-            open={mergedOpen}
-            onToggleOpen={onToggleOpen}
-            searchValue={mergedSearchValue}
-            activeValue={activeValue}
-            onSearch={triggerSearch}
-            onSearchSubmit={onSearchSubmit}
-            onSelect={onInternalSelectionSelect}
-            tokenWithEnter={tokenWithEnter}
-          />
-        </SelectTrigger>
+        {selectorNode}
 
         {arrowNode}
         {clearNode}
diff --git a/tests/Select.test.tsx b/tests/Select.test.tsx
index 184e61807..def4206fa 100644
--- a/tests/Select.test.tsx
+++ b/tests/Select.test.tsx
@@ -801,6 +801,24 @@ describe('Select.Basic', () => {
     expect(onCompositionEnd).toHaveBeenCalled();
   });
 
+  it('getRawInputElement for rc-cascader', () => {
+    const wrapper = mount(
+      <Select
+        getRawInputElement={() => <span className="bamboo" />}
+        options={[
+          {
+            label: <span className="little" />,
+            value: 'little',
+          },
+        ]}
+        open
+      />,
+    );
+
+    expect(wrapper.exists('.bamboo')).toBeTruthy();
+    expect(wrapper.exists('.little')).toBeTruthy();
+  });
+
   describe('propTypes', () => {
     let errorSpy;
 
diff --git a/typings/index.d.ts b/typings/index.d.ts
index f58471bdb..8c47d3f76 100644
--- a/typings/index.d.ts
+++ b/typings/index.d.ts
@@ -7,5 +7,3 @@ declare module 'rc-menu';
 declare module 'rc-util/lib/Children/toArray';
 
 declare module 'dom-scroll-into-view';
-
-declare module 'rc-trigger';