Refactor Badge and support count as custom component

close #11134
close #12140

components/badge/ScrollNumber.tsx

@@ -15,6 +15,7 @@ export interface ScrollNumberProps {
   prefixCls?: string;
   className?: string;
   count?: string | number | null;
+  displayComponent?: React.ReactElement<any>;
   component?: string;
   onAnimated?: Function;
   style?: React.CSSProperties;
@@ -116,22 +117,23 @@ export default class ScrollNumber extends Component<ScrollNumberProps, ScrollNum
   }
 
   renderNumberElement() {
-    const state = this.state;
-    if (!state.count || isNaN(state.count as number)) {
-      return state.count;
+    const { count } = this.state;
+    if (!count || isNaN(count as number)) {
+      return count;
     }
-    return getNumberArray(state.count)
+    return getNumberArray(count)
       .map((num, i) => this.renderCurrentNumber(num, i)).reverse();
   }
 
   render() {
-    const { prefixCls, className, style, title, component = 'sup' } = this.props;
+    const { prefixCls, className, style, title, component = 'sup', displayComponent } = this.props;
     // fix https://fb.me/react-unknown-prop
     const restProps = omit(this.props, [
       'count',
       'onAnimated',
       'component',
       'prefixCls',
+      'displayComponent',
     ]);
     const newProps = {
       ...restProps,
@@ -144,6 +146,11 @@ export default class ScrollNumber extends Component<ScrollNumberProps, ScrollNum
     if (style && style.borderColor) {
       newProps.style.boxShadow = `0 0 0 1px ${style.borderColor} inset`;
     }
+    if (displayComponent) {
+      return React.cloneElement(displayComponent, {
+        className: `${prefixCls}-custom-component`,
+      });
+    }
     return createElement(
       component as any,
       newProps,

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

@@ -186,6 +186,35 @@ exports[`renders ./components/badge/demo/basic.md correctly 1`] = `
       0
     </sup>
   </span>
+  <span
+    class="ant-badge"
+  >
+    <a
+      class="head-example"
+      href="#"
+    />
+    <i
+      class="anticon anticon-clock-circle ant-scroll-number-custom-component"
+      style="color:#f5222d"
+    >
+      <svg
+        aria-hidden="true"
+        class=""
+        data-icon="clock-circle"
+        fill="currentColor"
+        height="1em"
+        viewBox="64 64 896 896"
+        width="1em"
+      >
+        <path
+          d="M512 64C264.6 64 64 264.6 64 512s200.6 448 448 448 448-200.6 448-448S759.4 64 512 64zm0 820c-205.4 0-372-166.6-372-372s166.6-372 372-372 372 166.6 372 372-166.6 372-372 372z"
+        />
+        <path
+          d="M686.7 638.6L544.1 535.5V288c0-4.4-3.6-8-8-8H488c-4.4 0-8 3.6-8 8v275.4c0 2.6 1.2 5 3.3 6.5l165.4 120.6c3.6 2.6 8.6 1.8 11.2-1.7l28.6-39c2.6-3.7 1.8-8.7-1.8-11.2z"
+        />
+      </svg>
+    </i>
+  </span>
 </div>
 `;
 

components/badge/demo/basic.md

@@ -14,7 +14,7 @@ title:
 Simplest Usage. Badge will be hidden when `count` is `0`, but we can use `showZero` to show it.
 
 ````jsx
-import { Badge } from 'antd';
+import { Badge, Icon } from 'antd';
 
 ReactDOM.render(
   <div>
@@ -24,6 +24,9 @@ ReactDOM.render(
     <Badge count={0} showZero>
       <a href="#" className="head-example" />
     </Badge>
+    <Badge count={<Icon type="clock-circle" style={{ color: '#f5222d' }} />}>
+      <a href="#" className="head-example" />
+    </Badge>
   </div>,
   mountNode);
 ````

components/badge/index.tsx

@@ -8,7 +8,7 @@ export { ScrollNumberProps } from './ScrollNumber';
 
 export interface BadgeProps {
   /** Number to show in badge */
-  count?: number | string | null;
+  count?: React.ReactNode;
   showZero?: boolean;
   /** Max count to show */
   overflowCount?: number;
@@ -35,15 +35,125 @@ export default class Badge extends React.Component<BadgeProps, any> {
   };
 
   static propTypes = {
-    count: PropTypes.oneOfType([
-      PropTypes.string,
-      PropTypes.number,
-    ]),
+    count: PropTypes.node,
     showZero: PropTypes.bool,
     dot: PropTypes.bool,
     overflowCount: PropTypes.number,
   };
 
+  getBadgeClassName() {
+    const {
+      prefixCls,
+      className,
+      status,
+      children,
+    } = this.props;
+    return classNames(className, prefixCls, {
+      [`${prefixCls}-status`]: !!status,
+      [`${prefixCls}-not-a-wrapper`]: !children,
+    }) as string;
+  }
+
+  isZero() {
+    const numberedDispayCount = this.getNumberedDispayCount();
+    return numberedDispayCount === '0' || numberedDispayCount === 0;
+  }
+
+  isDot() {
+    const { dot, status } = this.props;
+    const isZero = this.isZero();
+    return (dot && !isZero) || status;
+  }
+
+  isHidden() {
+    const { showZero } = this.props;
+    const displayCount = this.getDispayCount();
+    const isZero = this.isZero();
+    const isDot = this.isDot();
+    const isEmpty = displayCount === null || displayCount === undefined || displayCount === '';
+    return (isEmpty || (isZero && !showZero)) && !isDot;
+  }
+
+  getNumberedDispayCount() {
+    const { count, overflowCount } = this.props;
+    const displayCount = (count as number) > (overflowCount as number) ? `${overflowCount}+` : count;
+    return displayCount as string | number | null;
+  }
+
+  getDispayCount() {
+    const isDot = this.isDot();
+    // dot mode don't need count
+    if (isDot) {
+      return '';
+    }
+    return this.getNumberedDispayCount();
+  }
+
+  getScollNumberTitle() {
+    const { title, count } = this.props;
+    if (title) {
+      return title;
+    }
+    return (typeof count === 'string' || typeof count === 'number') ? count : undefined;
+  }
+
+  getStyleWithOffset() {
+    const { offset, style } = this.props;
+    return offset ? {
+      right: -parseInt(offset[0] as string, 10),
+      marginTop: offset[1],
+      ...style,
+    } : style;
+  }
+
+  renderStatusText() {
+    const { prefixCls, text } = this.props;
+    const hidden = this.isHidden();
+    return (hidden || !text) ? null : (
+      <span className={`${prefixCls}-status-text`}>{text}</span>
+    );
+  }
+
+  renderDispayComponent() {
+    const { count } = this.props;
+    return (count && typeof count === 'object') ? (count as React.ReactElement<any>) : undefined;
+  }
+
+  renderBadgeNumber() {
+    const {
+      count,
+      prefixCls,
+      scrollNumberPrefixCls,
+      status,
+    } = this.props;
+
+    const displayCount = this.getDispayCount();
+    const isDot = this.isDot();
+    const hidden = this.isHidden();
+
+    const scrollNumberCls = classNames({
+      [`${prefixCls}-dot`]: isDot,
+      [`${prefixCls}-count`]: !isDot,
+      [`${prefixCls}-multiple-words`]: !isDot && count && count.toString && count.toString().length > 1,
+      [`${prefixCls}-status-${status}`]: !!status,
+    });
+
+    const styleWithOffset = this.getStyleWithOffset();
+
+    return hidden ? null : (
+      <ScrollNumber
+        prefixCls={scrollNumberPrefixCls}
+        data-show={!hidden}
+        className={scrollNumberCls}
+        count={displayCount}
+        displayComponent={this.renderDispayComponent()} // <Badge status="success" count={<Icon type="xxx" />}></Badge>
+        title={this.getScollNumberTitle()}
+        style={styleWithOffset}
+        key="scrollNumber"
+      />
+    );
+  }
+
   render() {
     const {
       count,
@@ -61,62 +171,29 @@ export default class Badge extends React.Component<BadgeProps, any> {
       title,
       ...restProps
     } = this.props;
-    let displayCount = (count as number) > (overflowCount as number) ? `${overflowCount}+` : count;
-    const isZero = displayCount === '0' || displayCount === 0;
-    const isDot = (dot && !isZero) || status;
-    // dot mode don't need count
-    if (isDot) {
-      displayCount = '';
-    }
-    const isEmpty = displayCount === null || displayCount === undefined || displayCount === '';
-    const hidden = (isEmpty || (isZero && !showZero)) && !isDot;
+
+    const scrollNumber = this.renderBadgeNumber();
+    const statusText = this.renderStatusText();
+
     const statusCls = classNames({
       [`${prefixCls}-status-dot`]: !!status,
       [`${prefixCls}-status-${status}`]: !!status,
     });
-    const scrollNumberCls = classNames({
-      [`${prefixCls}-dot`]: isDot,
-      [`${prefixCls}-count`]: !isDot,
-      [`${prefixCls}-multiple-words`]: !isDot && count && count.toString && count.toString().length > 1,
-      [`${prefixCls}-status-${status}`]: !!status,
-    });
-    const badgeCls = classNames(className, prefixCls, {
-      [`${prefixCls}-status`]: !!status,
-      [`${prefixCls}-not-a-wrapper`]: !children,
-    });
-    const styleWithOffset = offset ? {
-      right: -parseInt(offset[0] as string, 10),
-      marginTop: offset[1],
-      ...style,
-    } : style;
+
+    const styleWithOffset = this.getStyleWithOffset();
+
     // <Badge status="success" />
     if (!children && status) {
       return (
-        <span {...restProps} className={badgeCls} style={styleWithOffset}>
+        <span {...restProps} className={this.getBadgeClassName()} style={styleWithOffset}>
           <span className={statusCls} />
           <span className={`${prefixCls}-status-text`}>{text}</span>
         </span>
       );
     }
 
-    const scrollNumber = hidden ? null : (
-      <ScrollNumber
-        prefixCls={scrollNumberPrefixCls}
-        data-show={!hidden}
-        className={scrollNumberCls}
-        count={displayCount}
-        title={title || count}
-        style={styleWithOffset}
-        key="scrollNumber"
-      />
-    );
-
-    const statusText = (hidden || !text) ? null : (
-      <span className={`${prefixCls}-status-text`}>{text}</span>
-    );
-
     return (
-      <span {...restProps} className={badgeCls}>
+      <span {...restProps} className={this.getBadgeClassName()}>
         {children}
         <Animate
           component=""

components/badge/style/index.less

@@ -46,13 +46,18 @@
   }
 
   &-count,
-  &-dot {
+  &-dot,
+  .@{number-prefix-cls}-custom-component {
     position: absolute;
     right: 0;
     transform: translateX(50%);
     transform-origin: 100%;
   }
 
+  .@{number-prefix-cls}-custom-component {
+    transform: translate(50%, -50%);
+  }
+
   &-status {
     line-height: inherit;
     vertical-align: baseline;