refactor: wrap picker with FC (#34994)

* refactor: wrap picker with FC * chore: code clean * feat: cast type * test: fix test
ant-design · Apr 15, 2022 · 9b73058 · 9b73058
1 parent 7769ca1
commit 9b73058
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 15 deletions.
diff --git a/components/date-picker/generatePicker/generateRangePicker.tsx b/components/date-picker/generatePicker/generateRangePicker.tsx
@@ -1,4 +1,5 @@
 import * as React from 'react';
+import { forwardRef, useContext } from 'react';
 import classNames from 'classnames';
 import CalendarOutlined from '@ant-design/icons/CalendarOutlined';
 import ClockCircleOutlined from '@ant-design/icons/ClockCircleOutlined';
@@ -7,14 +8,14 @@ import SwapRightOutlined from '@ant-design/icons/SwapRightOutlined';
 import { RangePicker as RCRangePicker } from 'rc-picker';
 import { GenerateConfig } from 'rc-picker/lib/generate/index';
 import enUS from '../locale/en_US';
-import { ConfigContext, ConfigConsumerProps } from '../../config-provider';
+import { ConfigConsumerProps, ConfigContext } from '../../config-provider';
 import SizeContext from '../../config-provider/SizeContext';
 import LocaleReceiver from '../../locale-provider/LocaleReceiver';
 import { getRangePlaceholder, transPlacement2DropdownAlign } from '../util';
-import { RangePickerProps, PickerLocale, getTimeProps, Components } from '.';
-import { PickerComponentClass } from './interface';
+import { Components, getTimeProps, PickerLocale, RangePickerProps } from '.';
 import { FormItemInputContext } from '../../form/context';
 import { getMergedStatus, getStatusClassNames } from '../../_util/statusUtils';
+import { PickerComponentClass } from './interface';
 
 export default function generateRangePicker<DateType>(
   generateConfig: GenerateConfig<DateType>,
@@ -42,7 +43,7 @@ export default function generateRangePicker<DateType>(
       const locale = { ...contextLocale, ...this.props.locale };
       const { getPrefixCls, direction, getPopupContainer } = this.context;
       const {
-        prefixCls: customizePrefixCls,
+        prefixCls,
         getPopupContainer: customGetPopupContainer,
         className,
         placement,
@@ -53,7 +54,6 @@ export default function generateRangePicker<DateType>(
         ...restProps
       } = this.props;
       const { format, showTime, picker } = this.props as any;
-      const prefixCls = getPrefixCls('picker', customizePrefixCls);
 
       let additionalOverrideProps: any = {};
 
@@ -105,7 +105,7 @@ export default function generateRangePicker<DateType>(
                           [`${prefixCls}-borderless`]: !bordered,
                         },
                         getStatusClassNames(
-                          prefixCls,
+                          prefixCls as string,
                           getMergedStatus(contextStatus, customStatus),
                           hasFeedback,
                         ),
@@ -136,5 +136,12 @@ export default function generateRangePicker<DateType>(
     }
   }
 
-  return RangePicker;
+  return forwardRef<RangePicker, RangePickerProps<DateType>>((props, ref) => {
+    const { prefixCls: customizePrefixCls } = props;
+
+    const { getPrefixCls } = useContext(ConfigContext);
+    const prefixCls = getPrefixCls('picker', customizePrefixCls);
+
+    return <RangePicker {...props} prefixCls={prefixCls} ref={ref} />;
+  }) as unknown as PickerComponentClass<RangePickerProps<DateType>>;
 }
diff --git a/components/date-picker/generatePicker/generateSinglePicker.tsx b/components/date-picker/generatePicker/generateSinglePicker.tsx
@@ -6,6 +6,7 @@ import CloseCircleFilled from '@ant-design/icons/CloseCircleFilled';
 import RCPicker from 'rc-picker';
 import { PickerMode } from 'rc-picker/lib/interface';
 import { GenerateConfig } from 'rc-picker/lib/generate/index';
+import { forwardRef, useContext } from 'react';
 import enUS from '../locale/en_US';
 import { getPlaceholder, transPlacement2DropdownAlign } from '../util';
 import devWarning from '../../_util/devWarning';
@@ -20,9 +21,9 @@ import {
   getTimeProps,
   Components,
 } from '.';
-import { PickerComponentClass } from './interface';
 import { FormItemInputContext } from '../../form/context';
 import { getMergedStatus, getStatusClassNames, InputStatus } from '../../_util/statusUtils';
+import { DatePickRef, PickerComponentClass } from './interface';
 
 export default function generatePicker<DateType>(generateConfig: GenerateConfig<DateType>) {
   type DatePickerProps = PickerProps<DateType> & {
@@ -67,7 +68,7 @@ export default function generatePicker<DateType>(generateConfig: GenerateConfig<
         const locale = { ...contextLocale, ...this.props.locale };
         const { getPrefixCls, direction, getPopupContainer } = this.context;
         const {
-          prefixCls: customizePrefixCls,
+          prefixCls,
           getPopupContainer: customizeGetPopupContainer,
           className,
           size: customizeSize,
@@ -78,7 +79,6 @@ export default function generatePicker<DateType>(generateConfig: GenerateConfig<
           ...restProps
         } = this.props;
         const { format, showTime } = this.props as any;
-        const prefixCls = getPrefixCls('picker', customizePrefixCls);
 
         const additionalProps = {
           showToday: true,
@@ -137,7 +137,7 @@ export default function generatePicker<DateType>(generateConfig: GenerateConfig<
                             [`${prefixCls}-borderless`]: !bordered,
                           },
                           getStatusClassNames(
-                            prefixCls,
+                            prefixCls as string,
                             getMergedStatus(contextStatus, customStatus),
                             hasFeedback,
                           ),
@@ -167,11 +167,26 @@ export default function generatePicker<DateType>(generateConfig: GenerateConfig<
       }
     }
 
+    const PickerWrapper = forwardRef<DatePickRef<DateType>, InnerPickerProps>((props, ref) => {
+      const { prefixCls: customizePrefixCls } = props;
+
+      const { getPrefixCls } = useContext(ConfigContext);
+      const prefixCls = getPrefixCls('picker', customizePrefixCls);
+
+      const pickerProps: InnerPickerProps = {
+        ...props,
+        prefixCls,
+        ref,
+      };
+
+      return <Picker {...pickerProps} />;
+    });
+
     if (displayName) {
-      Picker.displayName = displayName;
+      PickerWrapper.displayName = displayName;
     }
 
-    return Picker as PickerComponentClass<InnerPickerProps>;
+    return PickerWrapper as unknown as PickerComponentClass<InnerPickerProps>;
   }
 
   const DatePicker = getPicker<DatePickerProps>();

diff --git a/components/time-picker/__tests__/index.test.js b/components/time-picker/__tests__/index.test.js
@@ -47,7 +47,7 @@ describe('TimePicker', () => {
   it('clearIcon should render correctly', () => {
     const clearIcon = <div className="test-clear-icon">test</div>;
     const wrapper = mount(<TimePicker clearIcon={clearIcon} />);
-    expect(wrapper.find('Picker').prop('clearIcon')).toEqual(
+    expect(wrapper.find('Picker').last().prop('clearIcon')).toEqual(
       <div className="test-clear-icon">test</div>,
     );
   });
@@ -70,7 +70,7 @@ describe('TimePicker', () => {
         popupClassName={popupClassName}
       />,
     );
-    expect(wrapper.find('Picker').prop('dropdownClassName')).toEqual(popupClassName);
+    expect(wrapper.find('Picker').last().prop('dropdownClassName')).toEqual(popupClassName);
   });
 
   it('should pass popupClassName prop to RangePicker as dropdownClassName prop', () => {