refactor(badge): 🛠 improve code (#24601)

* refactor(badge): improve code

* refactor(badge): improve code

* increase coverage

components/badge/ScrollNumber.tsx

@@ -1,7 +1,6 @@
 import * as React from 'react';
-import omit from 'omit.js';
 import classNames from 'classnames';
-import { ConfigConsumer, ConfigConsumerProps } from '../config-provider';
+import { ConfigContext } from '../config-provider';
 import { cloneElement } from '../_util/reactNode';
 
 function getNumberArray(num: string | number | undefined | null) {
@@ -51,15 +50,27 @@ export interface ScrollNumberState {
   count?: string | number | null;
 }
 
-const ScrollNumber: React.FC<ScrollNumberProps> = props => {
+const ScrollNumber: React.FC<ScrollNumberProps> = ({
+  prefixCls: customizePrefixCls,
+  count: customizeCount,
+  className,
+  style,
+  title,
+  component = 'sup',
+  displayComponent,
+  onAnimated = () => {},
+  ...restProps
+}) => {
   const [animateStarted, setAnimateStarted] = React.useState(true);
-  const [count, setCount] = React.useState(props.count);
+  const [count, setCount] = React.useState(customizeCount);
-  const [prevCount, setPrevCount] = React.useState(props.count);
+  const [prevCount, setPrevCount] = React.useState(customizeCount);
-  const [lastCount, setLastCount] = React.useState(props.count);
+  const [lastCount, setLastCount] = React.useState(customizeCount);
+  const { getPrefixCls } = React.useContext(ConfigContext);
+  const prefixCls = getPrefixCls('scroll-number', customizePrefixCls);
 
-  if (prevCount !== props.count) {
+  if (prevCount !== customizeCount) {
     setAnimateStarted(true);
-    setPrevCount(props.count);
+    setPrevCount(customizeCount);
   }
 
   React.useEffect(() => {
@@ -70,10 +81,8 @@ const ScrollNumber: React.FC<ScrollNumberProps> = props => {
       // performing the transition.
       timeout = setTimeout(() => {
         setAnimateStarted(false);
-        setCount(props.count);
+        setCount(customizeCount);
-        if (props.onAnimated) {
+        onAnimated();
-          props.onAnimated();
-        }
       });
     }
     return () => {
@@ -81,7 +90,7 @@ const ScrollNumber: React.FC<ScrollNumberProps> = props => {
         clearTimeout(timeout);
       }
     };
-  }, [animateStarted, count, props.count, props.onAnimated]);
+  }, [animateStarted, customizeCount, onAnimated]);
 
   const getPositionByNum = (num: number, i: number) => {
     const currentCount = Math.abs(Number(count));
@@ -106,7 +115,7 @@ const ScrollNumber: React.FC<ScrollNumberProps> = props => {
     return num;
   };
 
-  const renderCurrentNumber = (prefixCls: string, num: number | string, i: number) => {
+  const renderCurrentNumber = (num: number | string, i: number) => {
     if (typeof num === 'number') {
       const position = getPositionByNum(num, i);
       const removeTransition = animateStarted || getNumberArray(lastCount)[i] === undefined;
@@ -133,65 +142,40 @@ const ScrollNumber: React.FC<ScrollNumberProps> = props => {
     );
   };
 
-  const renderNumberElement = (prefixCls: string) => {
+  const renderNumberElement = () => {
     if (count && Number(count) % 1 === 0) {
       return getNumberArray(count)
-        .map((num, i) => renderCurrentNumber(prefixCls, num, i))
+        .map((num, i) => renderCurrentNumber(num, i))
         .reverse();
     }
     return count;
   };
 
-  const renderScrollNumber = ({ getPrefixCls }: ConfigConsumerProps) => {
+  const newProps = {
-    const {
+    ...restProps,
-      prefixCls: customizePrefixCls,
+    style,
-      className,
+    className: classNames(prefixCls, className),
-      style,
+    title: title as string,
-      title,
-      component = 'sup',
-      displayComponent,
-    } = props;
-    // fix https://fb.me/react-unknown-prop
-    const restProps = omit(props, [
-      'count',
-      'onAnimated',
-      'component',
-      'prefixCls',
-      'displayComponent',
-    ]);
-    const prefixCls = getPrefixCls('scroll-number', customizePrefixCls);
-    const newProps = {
-      ...restProps,
-      className: classNames(prefixCls, className),
-      title: title as string,
-    };
-
-    // allow specify the border
-    // mock border-color by box-shadow for compatible with old usage:
-    // <Badge count={4} style={{ backgroundColor: '#fff', color: '#999', borderColor: '#d9d9d9' }} />
-    if (style && style.borderColor) {
-      newProps.style = {
-        ...style,
-        boxShadow: `0 0 0 1px ${style.borderColor} inset`,
-      };
-    }
-    if (displayComponent) {
-      return cloneElement(displayComponent, {
-        className: classNames(
-          `${prefixCls}-custom-component`,
-          displayComponent.props && displayComponent.props.className,
-        ),
-      });
-    }
-    return React.createElement(component as any, newProps, renderNumberElement(prefixCls));
   };
 
-  return <ConfigConsumer>{renderScrollNumber}</ConfigConsumer>;
+  // allow specify the border
-};
+  // mock border-color by box-shadow for compatible with old usage:
-
+  // <Badge count={4} style={{ backgroundColor: '#fff', color: '#999', borderColor: '#d9d9d9' }} />
-ScrollNumber.defaultProps = {
+  if (style && style.borderColor) {
-  count: null,
+    newProps.style = {
-  onAnimated() {},
+      ...style,
+      boxShadow: `0 0 0 1px ${style.borderColor} inset`,
+    };
+  }
+  if (displayComponent) {
+    return cloneElement(displayComponent, {
+      className: classNames(
+        `${prefixCls}-custom-component`,
+        displayComponent.props && displayComponent.props.className,
+      ),
+    });
+  }
+  return React.createElement(component as any, newProps, renderNumberElement());
 };
 
 export default ScrollNumber;

components/badge/__tests__/__snapshots__/index.test.js.snap

@@ -17,9 +17,6 @@ exports[`Badge badge should support float number 1`] = `
 exports[`Badge badge should support float number 2`] = `
 <Badge
   count="3.5"
-  dot={false}
-  overflowCount={99}
-  showZero={false}
 >
   <span
     className="ant-badge ant-badge-not-a-wrapper"
@@ -51,7 +48,6 @@ exports[`Badge badge should support float number 2`] = `
           count="3.5"
           data-show={true}
           key="scrollNumber"
-          onAnimated={[Function]}
           prefixCls="ant-scroll-number"
           title="3.5"
         >
@@ -938,9 +934,6 @@ exports[`Badge should be compatible with borderColor style 1`] = `
 exports[`Badge should render when count is changed 1`] = `
 <Badge
   count={10}
-  dot={false}
-  overflowCount={99}
-  showZero={false}
 >
   <span
     className="ant-badge ant-badge-not-a-wrapper"
@@ -972,7 +965,6 @@ exports[`Badge should render when count is changed 1`] = `
           count={10}
           data-show={true}
           key="scrollNumber"
-          onAnimated={[Function]}
           prefixCls="ant-scroll-number"
           title={10}
         >
@@ -1185,9 +1177,6 @@ exports[`Badge should render when count is changed 1`] = `
 exports[`Badge should render when count is changed 2`] = `
 <Badge
   count={11}
-  dot={false}
-  overflowCount={99}
-  showZero={false}
 >
   <span
     className="ant-badge ant-badge-not-a-wrapper"
@@ -1219,7 +1208,6 @@ exports[`Badge should render when count is changed 2`] = `
           count={11}
           data-show={true}
           key="scrollNumber"
-          onAnimated={[Function]}
           prefixCls="ant-scroll-number"
           title={11}
         >
@@ -1625,9 +1613,6 @@ exports[`Badge should render when count is changed 2`] = `
 exports[`Badge should render when count is changed 3`] = `
 <Badge
   count={11}
-  dot={false}
-  overflowCount={99}
-  showZero={false}
 >
   <span
     className="ant-badge ant-badge-not-a-wrapper"
@@ -1659,7 +1644,6 @@ exports[`Badge should render when count is changed 3`] = `
           count={11}
           data-show={true}
           key="scrollNumber"
-          onAnimated={[Function]}
           prefixCls="ant-scroll-number"
           title={11}
         >
@@ -2064,10 +2048,7 @@ exports[`Badge should render when count is changed 3`] = `
 
 exports[`Badge should render when count is changed 4`] = `
 <Badge
-  count={10}
+  count={111}
-  dot={false}
-  overflowCount={99}
-  showZero={false}
 >
   <span
     className="ant-badge ant-badge-not-a-wrapper"
@@ -2096,17 +2077,16 @@ exports[`Badge should render when count is changed 4`] = `
       >
         <ScrollNumber
           className="ant-badge-count ant-badge-multiple-words"
-          count={10}
+          count="99+"
           data-show={true}
           key="scrollNumber"
-          onAnimated={[Function]}
           prefixCls="ant-scroll-number"
-          title={10}
+          title={111}
         >
           <sup
             className="ant-scroll-number ant-badge-count ant-badge-multiple-words"
             data-show={true}
-            title={10}
+            title={111}
           >
             <span
               className="ant-scroll-number-only"
@@ -2503,11 +2483,59 @@ exports[`Badge should render when count is changed 4`] = `
 `;
 
 exports[`Badge should render when count is changed 5`] = `
+<Badge
+  count={10}
+>
+  <span
+    className="ant-badge ant-badge-not-a-wrapper"
+  >
+    <Animate
+      animation={Object {}}
+      component=""
+      componentProps={Object {}}
+      onAppear={[Function]}
+      onEnd={[Function]}
+      onEnter={[Function]}
+      onLeave={[Function]}
+      showProp="data-show"
+      transitionAppear={true}
+      transitionEnter={true}
+      transitionLeave={true}
+      transitionName=""
+    >
+      <AnimateChild
+        animation={Object {}}
+        key="scrollNumber"
+        transitionAppear={true}
+        transitionEnter={true}
+        transitionLeave={true}
+        transitionName=""
+      >
+        <ScrollNumber
+          className="ant-badge-count ant-badge-multiple-words"
+          count={10}
+          data-show={true}
+          key="scrollNumber"
+          prefixCls="ant-scroll-number"
+          title={10}
+        >
+          <sup
+            className="ant-scroll-number ant-badge-count ant-badge-multiple-words"
+            data-show={true}
+            title={10}
+          >
+            99+
+          </sup>
+        </ScrollNumber>
+      </AnimateChild>
+    </Animate>
+  </span>
+</Badge>
+`;
+
+exports[`Badge should render when count is changed 6`] = `
 <Badge
   count={9}
-  dot={false}
-  overflowCount={99}
-  showZero={false}
 >
   <span
     className="ant-badge ant-badge-not-a-wrapper"
@@ -2539,7 +2567,6 @@ exports[`Badge should render when count is changed 5`] = `
           count={9}
           data-show={true}
           key="scrollNumber"
-          onAnimated={[Function]}
           prefixCls="ant-scroll-number"
           title={9}
         >

components/badge/__tests__/index.test.js

@@ -66,6 +66,9 @@ describe('Badge', () => {
     wrapper.setProps({ count: 11 });
     jest.runAllTimers();
     expect(wrapper).toMatchSnapshot();
+    wrapper.setProps({ count: 111 });
+    jest.runAllTimers();
+    expect(wrapper).toMatchSnapshot();
     wrapper.setProps({ count: 10 });
     jest.runAllTimers();
     expect(wrapper).toMatchSnapshot();

components/badge/index.tsx

@@ -1,10 +1,9 @@
 import * as React from 'react';
 import Animate from 'rc-animate';
-import omit from 'omit.js';
 import classNames from 'classnames';
 import ScrollNumber from './ScrollNumber';
 import { PresetColorTypes, PresetColorType, PresetStatusColorType } from '../_util/colors';
-import { ConfigConsumer, ConfigConsumerProps } from '../config-provider';
+import { ConfigContext } from '../config-provider';
 import { LiteralUnion } from '../_util/type';
 import { cloneElement } from '../_util/reactNode';
 
@@ -33,16 +32,33 @@ function isPresetColor(color?: string): boolean {
   return (PresetColorTypes as any[]).indexOf(color) !== -1;
 }
 
-const Badge: React.FC<BadgeProps> = props => {
+const Badge: React.FC<BadgeProps> = ({
+  prefixCls: customizePrefixCls,
+  scrollNumberPrefixCls: customizeScrollNumberPrefixCls,
+  children,
+  status,
+  text,
+  color,
+  count = null,
+  overflowCount = 99,
+  dot = false,
+  title,
+  offset,
+  style,
+  className,
+  showZero = false,
+  ...restProps
+}) => {
+  const { getPrefixCls, direction } = React.useContext(ConfigContext);
+  const prefixCls = getPrefixCls('badge', customizePrefixCls);
+
   const getNumberedDisplayCount = () => {
-    const { count, overflowCount } = props;
     const displayCount =
       (count as number) > (overflowCount as number) ? `${overflowCount}+` : count;
     return displayCount as string | number | null;
   };
 
   const hasStatus = (): boolean => {
-    const { status, color } = props;
     return !!status || !!color;
   };
 
@@ -52,7 +68,6 @@ const Badge: React.FC<BadgeProps> = props => {
   };
 
   const isDot = () => {
-    const { dot } = props;
     return (dot && !isZero()) || hasStatus();
   };
 
@@ -65,7 +80,6 @@ const Badge: React.FC<BadgeProps> = props => {
   };
 
   const getScrollNumberTitle = () => {
-    const { title, count } = props;
     if (title) {
       return title;
     }
@@ -73,7 +87,6 @@ const Badge: React.FC<BadgeProps> = props => {
   };
 
   const getStyleWithOffset = () => {
-    const { offset, style } = props;
     return offset
       ? {
           right: -parseInt(offset[0] as string, 10),
@@ -83,30 +96,18 @@ const Badge: React.FC<BadgeProps> = props => {
       : style;
   };
 
-  const getBadgeClassName = (prefixCls: string, direction: string = 'ltr') => {
-    const { className, children } = props;
-    return classNames(className, prefixCls, {
-      [`${prefixCls}-status`]: hasStatus(),
-      [`${prefixCls}-not-a-wrapper`]: !children,
-      [`${prefixCls}-rtl`]: direction === 'rtl',
-    }) as string;
-  };
-
   const isHidden = () => {
-    const { showZero } = props;
     const displayCount = getDisplayCount();
     const isEmpty = displayCount === null || displayCount === undefined || displayCount === '';
     return (isEmpty || (isZero() && !showZero)) && !isDot();
   };
 
-  const renderStatusText = (prefixCls: string) => {
+  const renderStatusText = () => {
-    const { text } = props;
     const hidden = isHidden();
     return hidden || !text ? null : <span className={`${prefixCls}-status-text`}>{text}</span>;
   };
 
   const renderDisplayComponent = () => {
-    const { count } = props;
     const customNode = count as React.ReactElement<any>;
     if (!customNode || typeof customNode !== 'object') {
       return undefined;
@@ -119,18 +120,17 @@ const Badge: React.FC<BadgeProps> = props => {
     });
   };
 
-  const renderBadgeNumber = (prefixCls: string, scrollNumberPrefixCls: string) => {
+  const renderBadgeNumber = () => {
-    const { status, count, color } = props;
+    const scrollNumberPrefixCls = getPrefixCls('scroll-number', customizeScrollNumberPrefixCls);
-
     const displayCount = getDisplayCount();
-    const dot = isDot();
+    const bDot = isDot();
     const hidden = isHidden();
 
     const scrollNumberCls = classNames({
-      [`${prefixCls}-dot`]: dot,
+      [`${prefixCls}-dot`]: bDot,
-      [`${prefixCls}-count`]: !dot,
+      [`${prefixCls}-count`]: !bDot,
       [`${prefixCls}-multiple-words`]:
-        !dot && count && count.toString && count.toString().length > 1,
+        !bDot && count && count.toString && count.toString().length > 1,
       [`${prefixCls}-status-${status}`]: !!status,
       [`${prefixCls}-status-${color}`]: isPresetColor(color),
     });
@@ -155,85 +155,50 @@ const Badge: React.FC<BadgeProps> = props => {
     );
   };
 
-  const renderBadge = ({ getPrefixCls, direction }: ConfigConsumerProps) => {
+  const statusCls = classNames({
-    const {
+    [`${prefixCls}-status-dot`]: hasStatus(),
-      prefixCls: customizePrefixCls,
+    [`${prefixCls}-status-${status}`]: !!status,
-      scrollNumberPrefixCls: customizeScrollNumberPrefixCls,
+    [`${prefixCls}-status-${color}`]: isPresetColor(color),
-      children,
+  });
-      status,
+  const statusStyle: React.CSSProperties = {};
-      text,
+  if (color && !isPresetColor(color)) {
-      color,
+    statusStyle.background = color;
-      ...restProps
+  }
-    } = props;
-    const omitArr = [
-      'count',
-      'showZero',
-      'overflowCount',
-      'className',
-      'style',
-      'dot',
-      'offset',
-      'title',
-    ];
 
-    const prefixCls = getPrefixCls('badge', customizePrefixCls);
+  const badgeClassName = classNames(className, prefixCls, {
-    const scrollNumberPrefixCls = getPrefixCls('scroll-number', customizeScrollNumberPrefixCls);
+    [`${prefixCls}-status`]: hasStatus(),
-
+    [`${prefixCls}-not-a-wrapper`]: !children,
-    const scrollNumber = renderBadgeNumber(prefixCls, scrollNumberPrefixCls);
+    [`${prefixCls}-rtl`]: direction === 'rtl',
-    const statusText = renderStatusText(prefixCls);
+  });
-
-    const statusCls = classNames({
-      [`${prefixCls}-status-dot`]: hasStatus(),
-      [`${prefixCls}-status-${status}`]: !!status,
-      [`${prefixCls}-status-${color}`]: isPresetColor(color),
-    });
-    const statusStyle: React.CSSProperties = {};
-    if (color && !isPresetColor(color)) {
-      statusStyle.background = color;
-    }
-
-    // <Badge status="success" />
-    if (!children && hasStatus()) {
-      const styleWithOffset = getStyleWithOffset();
-      const statusTextColor = styleWithOffset && styleWithOffset.color;
-      return (
-        <span
-          {...omit(restProps, omitArr)}
-          className={getBadgeClassName(prefixCls, direction)}
-          style={styleWithOffset}
-        >
-          <span className={statusCls} style={statusStyle} />
-          <span style={{ color: statusTextColor }} className={`${prefixCls}-status-text`}>
-            {text}
-          </span>
-        </span>
-      );
-    }
 
+  // <Badge status="success" />
+  if (!children && hasStatus()) {
+    const styleWithOffset = getStyleWithOffset();
+    const statusTextColor = styleWithOffset && styleWithOffset.color;
     return (
-      <span {...omit(restProps, omitArr)} className={getBadgeClassName(prefixCls, direction)}>
+      <span {...restProps} className={badgeClassName} style={styleWithOffset}>
-        {children}
+        <span className={statusCls} style={statusStyle} />
-        <Animate
+        <span style={{ color: statusTextColor }} className={`${prefixCls}-status-text`}>
-          component=""
+          {text}
-          showProp="data-show"
+        </span>
-          transitionName={children ? `${prefixCls}-zoom` : ''}
-          transitionAppear
-        >
-          {scrollNumber}
-        </Animate>
-        {statusText}
       </span>
     );
-  };
+  }
 
-  return <ConfigConsumer>{renderBadge}</ConfigConsumer>;
+  return (
-};
+    <span {...restProps} className={badgeClassName}>
-
+      {children}
-Badge.defaultProps = {
+      <Animate
-  count: null,
+        component=""
-  showZero: false,
+        showProp="data-show"
-  dot: false,
+        transitionName={children ? `${prefixCls}-zoom` : ''}
-  overflowCount: 99,
+        transitionAppear
+      >
+        {renderBadgeNumber()}
+      </Animate>
+      {renderStatusText()}
+    </span>
+  );
 };
 
 export default Badge;