From da08be5d5596bf87f2a859207401eccaf9b902ba Mon Sep 17 00:00:00 2001
From: unix <unix.bio@gmail.com>
Date: Thu, 21 May 2020 16:45:36 +0800
Subject: [PATCH] feat(button): forward the ref of the original element

---
 .../__snapshots__/icon.test.tsx.snap          | 228 ++++++-------
 .../__snapshots__/index.test.tsx.snap         |  28 +-
 components/button/__tests__/index.test.tsx    |   7 +
 components/button/button.tsx                  | 307 ++++++++++--------
 4 files changed, 288 insertions(+), 282 deletions(-)

diff --git a/components/button/__tests__/__snapshots__/icon.test.tsx.snap b/components/button/__tests__/__snapshots__/icon.test.tsx.snap
index 19a2c49..55796dc 100644
--- a/components/button/__tests__/__snapshots__/icon.test.tsx.snap
+++ b/components/button/__tests__/__snapshots__/icon.test.tsx.snap
@@ -33,66 +33,66 @@ exports[`ButtonIcon should render correctly 1`] = `
             padding-right: 0;
           }
         </style></div><style>
-        .btn {
-          box-sizing: border-box;
-          display: inline-block;
-          padding: 0 1.375rem;
-          height: 2.5rem;
-          line-height: 2.5rem;
-          min-width: 12.5rem;
-          width: auto;
-          border-radius: 5px;
-          font-weight: 400;
-          font-size: .875rem;
-          user-select: none;
-          outline: none;
-          text-transform: capitalize;
-          justify-content: center;
-          text-align: center;
-          white-space: nowrap;
-          transition: background-color 200ms ease 0ms, box-shadow 200ms ease 0ms,
-            border 200ms ease 0ms, color 200ms ease 0ms;
-          position: relative;
-          overflow: hidden;
-          color: #666;
-          background-color: #fff;
-          border: 1px solid #eaeaea;
-          cursor: pointer;
-          pointer-events: auto;
-          box-shadow: none;
-          --zeit-ui-button-padding: 1.375rem;
-          --zeit-ui-button-height: 2.5rem;
-          --zeit-ui-button-color: #666;
-        }
+          .btn {
+            box-sizing: border-box;
+            display: inline-block;
+            padding: 0 1.375rem;
+            height: 2.5rem;
+            line-height: 2.5rem;
+            min-width: 12.5rem;
+            width: auto;
+            border-radius: 5px;
+            font-weight: 400;
+            font-size: .875rem;
+            user-select: none;
+            outline: none;
+            text-transform: capitalize;
+            justify-content: center;
+            text-align: center;
+            white-space: nowrap;
+            transition: background-color 200ms ease 0ms, box-shadow 200ms ease 0ms,
+              border 200ms ease 0ms, color 200ms ease 0ms;
+            position: relative;
+            overflow: hidden;
+            color: #666;
+            background-color: #fff;
+            border: 1px solid #eaeaea;
+            cursor: pointer;
+            pointer-events: auto;
+            box-shadow: none;
+            --zeit-ui-button-padding: 1.375rem;
+            --zeit-ui-button-height: 2.5rem;
+            --zeit-ui-button-color: #666;
+          }
 
-        .btn:hover {
-          color: #000;
-          --zeit-ui-button-color: #000;
-          background-color: #fff;
-          border-color: #000;
-          cursor: pointer;
-          pointer-events: auto;
-          box-shadow: none;
-          transform: translate3d(0px, 0px, 0px);
-        }
+          .btn:hover {
+            color: #000;
+            --zeit-ui-button-color: #000;
+            background-color: #fff;
+            border-color: #000;
+            cursor: pointer;
+            pointer-events: auto;
+            box-shadow: none;
+            transform: translate3d(0px, 0px, 0px);
+          }
 
-        .btn :global(.text) {
-          position: relative;
-          z-index: 1;
-          display: inline-flex;
-          justify-content: center;
-          align-items: center;
-          text-align: center;
-          line-height: inherit;
-          top: -1px;
-        }
+          .btn :global(.text) {
+            position: relative;
+            z-index: 1;
+            display: inline-flex;
+            justify-content: center;
+            align-items: center;
+            text-align: center;
+            line-height: inherit;
+            top: -1px;
+          }
 
-        .btn :global(.text p),
-        .btn :global(.text pre),
-        .btn :global(.text div) {
-          margin: 0;
-        }
-      </style></button>"
+          .btn :global(.text p),
+          .btn :global(.text pre),
+          .btn :global(.text div) {
+            margin: 0;
+          }
+        </style></button>"
 `;
 
 exports[`ButtonIcon should work with right 1`] = `
@@ -128,64 +128,64 @@ exports[`ButtonIcon should work with right 1`] = `
             padding-right: 0;
           }
         </style></div><style>
-        .btn {
-          box-sizing: border-box;
-          display: inline-block;
-          padding: 0 1.375rem;
-          height: 2.5rem;
-          line-height: 2.5rem;
-          min-width: 12.5rem;
-          width: auto;
-          border-radius: 5px;
-          font-weight: 400;
-          font-size: .875rem;
-          user-select: none;
-          outline: none;
-          text-transform: capitalize;
-          justify-content: center;
-          text-align: center;
-          white-space: nowrap;
-          transition: background-color 200ms ease 0ms, box-shadow 200ms ease 0ms,
-            border 200ms ease 0ms, color 200ms ease 0ms;
-          position: relative;
-          overflow: hidden;
-          color: #666;
-          background-color: #fff;
-          border: 1px solid #eaeaea;
-          cursor: pointer;
-          pointer-events: auto;
-          box-shadow: none;
-          --zeit-ui-button-padding: 1.375rem;
-          --zeit-ui-button-height: 2.5rem;
-          --zeit-ui-button-color: #666;
-        }
+          .btn {
+            box-sizing: border-box;
+            display: inline-block;
+            padding: 0 1.375rem;
+            height: 2.5rem;
+            line-height: 2.5rem;
+            min-width: 12.5rem;
+            width: auto;
+            border-radius: 5px;
+            font-weight: 400;
+            font-size: .875rem;
+            user-select: none;
+            outline: none;
+            text-transform: capitalize;
+            justify-content: center;
+            text-align: center;
+            white-space: nowrap;
+            transition: background-color 200ms ease 0ms, box-shadow 200ms ease 0ms,
+              border 200ms ease 0ms, color 200ms ease 0ms;
+            position: relative;
+            overflow: hidden;
+            color: #666;
+            background-color: #fff;
+            border: 1px solid #eaeaea;
+            cursor: pointer;
+            pointer-events: auto;
+            box-shadow: none;
+            --zeit-ui-button-padding: 1.375rem;
+            --zeit-ui-button-height: 2.5rem;
+            --zeit-ui-button-color: #666;
+          }
 
-        .btn:hover {
-          color: #000;
-          --zeit-ui-button-color: #000;
-          background-color: #fff;
-          border-color: #000;
-          cursor: pointer;
-          pointer-events: auto;
-          box-shadow: none;
-          transform: translate3d(0px, 0px, 0px);
-        }
+          .btn:hover {
+            color: #000;
+            --zeit-ui-button-color: #000;
+            background-color: #fff;
+            border-color: #000;
+            cursor: pointer;
+            pointer-events: auto;
+            box-shadow: none;
+            transform: translate3d(0px, 0px, 0px);
+          }
 
-        .btn :global(.text) {
-          position: relative;
-          z-index: 1;
-          display: inline-flex;
-          justify-content: center;
-          align-items: center;
-          text-align: center;
-          line-height: inherit;
-          top: -1px;
-        }
+          .btn :global(.text) {
+            position: relative;
+            z-index: 1;
+            display: inline-flex;
+            justify-content: center;
+            align-items: center;
+            text-align: center;
+            line-height: inherit;
+            top: -1px;
+          }
 
-        .btn :global(.text p),
-        .btn :global(.text pre),
-        .btn :global(.text div) {
-          margin: 0;
-        }
-      </style></button>"
+          .btn :global(.text p),
+          .btn :global(.text pre),
+          .btn :global(.text div) {
+            margin: 0;
+          }
+        </style></button>"
 `;
diff --git a/components/button/__tests__/__snapshots__/index.test.tsx.snap b/components/button/__tests__/__snapshots__/index.test.tsx.snap
index 3c5757d..470d7a9 100644
--- a/components/button/__tests__/__snapshots__/index.test.tsx.snap
+++ b/components/button/__tests__/__snapshots__/index.test.tsx.snap
@@ -1,35 +1,13 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`Button should render empty button correctly 1`] = `
-<Memo(Button)
-  auto={false}
-  className=""
-  disabled={false}
-  effect={true}
-  ghost={false}
-  htmlType="button"
-  loading={false}
-  shadow={false}
-  size="medium"
-  type="default"
-/>
-`;
+exports[`Button should render empty button correctly 1`] = `<Memo />`;
 
 exports[`Button should render special styles 1`] = `ReactWrapper {}`;
 
 exports[`Button should render special styles 2`] = `
-<Memo(Button)
-  auto={false}
-  className=""
-  disabled={false}
-  effect={true}
-  ghost={false}
-  htmlType="button"
+<Memo
   loading={true}
-  shadow={false}
-  size="medium"
-  type="default"
 >
   button
-</Memo(Button)>
+</Memo>
 `;
diff --git a/components/button/__tests__/index.test.tsx b/components/button/__tests__/index.test.tsx
index eb5434d..725786e 100644
--- a/components/button/__tests__/index.test.tsx
+++ b/components/button/__tests__/index.test.tsx
@@ -131,4 +131,11 @@ describe('Button', () => {
     await sleep(500)
     expect(wrapper.find('.loading-container').length).not.toBe(0)
   })
+
+  it('ref should be forwarded', () => {
+    const ref = React.createRef<HTMLButtonElement>()
+    const wrapper = mount(<Button ref={ref}>action</Button>)
+    expect(wrapper.find('button').getDOMNode()).toEqual(ref.current)
+    expect(() => wrapper.unmount()).not.toThrow()
+  })
 })
diff --git a/components/button/button.tsx b/components/button/button.tsx
index 51a4140..e091003 100644
--- a/components/button/button.tsx
+++ b/components/button/button.tsx
@@ -1,6 +1,13 @@
-import React, { useRef, useState, MouseEvent, useMemo } from 'react'
+import React, {
+  useRef,
+  useState,
+  MouseEvent,
+  useMemo,
+  useImperativeHandle,
+  PropsWithoutRef,
+  RefAttributes,
+} from 'react'
 import useTheme from '../styles/use-theme'
-import withDefaults from '../utils/with-defaults'
 import ButtonDrip from './button.drip'
 import ButtonLoading from '../loading'
 import { ButtonTypes, NormalSizes } from '../utils/prop-types'
@@ -46,157 +53,171 @@ const defaultProps = {
 type NativeAttrs = Omit<React.ButtonHTMLAttributes<any>, keyof Props>
 export type ButtonProps = Props & typeof defaultProps & NativeAttrs
 
-const Button: React.FC<React.PropsWithChildren<ButtonProps>> = ({ ...btnProps }) => {
-  const theme = useTheme()
-  const buttonRef = useRef<HTMLButtonElement>(null)
-  const [dripShow, setDripShow] = useState<boolean>(false)
-  const [dripX, setDripX] = useState<number>(0)
-  const [dripY, setDripY] = useState<number>(0)
-  const groupConfig = useButtonGroupContext()
-  const filteredProps = filterPropsWithGroup(btnProps, groupConfig)
-  const {
-    children,
-    disabled,
-    type,
-    loading,
-    shadow,
-    ghost,
-    effect,
-    onClick,
-    auto,
-    size,
-    icon,
-    htmlType,
-    iconRight,
-    className,
-    ...props
-  } = filteredProps
+const Button = React.forwardRef<HTMLButtonElement, React.PropsWithChildren<ButtonProps>>(
+  ({ ...btnProps }, ref: React.Ref<HTMLButtonElement | null>) => {
+    const theme = useTheme()
+    const buttonRef = useRef<HTMLButtonElement>(null)
+    useImperativeHandle(ref, () => buttonRef.current)
 
-  const { bg, border, color } = useMemo(() => getButtonColors(theme.palette, filteredProps), [
-    theme.palette,
-    filteredProps,
-  ])
-  const hover = useMemo(() => getButtonHoverColors(theme.palette, filteredProps), [
-    theme.palette,
-    filteredProps,
-  ])
-  const { cursor, events } = useMemo(() => getButtonCursor(disabled, loading), [disabled, loading])
-  const { height, minWidth, padding, width, fontSize } = useMemo(() => getButtonSize(size, auto), [
-    size,
-    auto,
-  ])
-  const dripColor = useMemo(() => getButtonDripColor(theme.palette, filteredProps), [
-    theme.palette,
-    filteredProps,
-  ])
+    const [dripShow, setDripShow] = useState<boolean>(false)
+    const [dripX, setDripX] = useState<number>(0)
+    const [dripY, setDripY] = useState<number>(0)
+    const groupConfig = useButtonGroupContext()
+    const filteredProps = filterPropsWithGroup(btnProps, groupConfig)
+    const {
+      children,
+      disabled,
+      type,
+      loading,
+      shadow,
+      ghost,
+      effect,
+      onClick,
+      auto,
+      size,
+      icon,
+      htmlType,
+      iconRight,
+      className,
+      ...props
+    } = filteredProps
 
-  /* istanbul ignore next */
-  const dripCompletedHandle = () => {
-    setDripShow(false)
-    setDripX(0)
-    setDripY(0)
-  }
+    const { bg, border, color } = useMemo(() => getButtonColors(theme.palette, filteredProps), [
+      theme.palette,
+      filteredProps,
+    ])
+    const hover = useMemo(() => getButtonHoverColors(theme.palette, filteredProps), [
+      theme.palette,
+      filteredProps,
+    ])
+    const { cursor, events } = useMemo(() => getButtonCursor(disabled, loading), [
+      disabled,
+      loading,
+    ])
+    const { height, minWidth, padding, width, fontSize } = useMemo(
+      () => getButtonSize(size, auto),
+      [size, auto],
+    )
+    const dripColor = useMemo(() => getButtonDripColor(theme.palette, filteredProps), [
+      theme.palette,
+      filteredProps,
+    ])
 
-  const clickHandler = (event: MouseEvent<HTMLButtonElement>) => {
-    if (disabled || loading) return
-    const showDrip = !shadow && !ghost && effect
     /* istanbul ignore next */
-    if (showDrip && buttonRef.current) {
-      const rect = buttonRef.current.getBoundingClientRect()
-      setDripShow(true)
-      setDripX(event.clientX - rect.left)
-      setDripY(event.clientY - rect.top)
+    const dripCompletedHandle = () => {
+      setDripShow(false)
+      setDripX(0)
+      setDripY(0)
     }
 
-    onClick && onClick(event)
-  }
+    const clickHandler = (event: MouseEvent<HTMLButtonElement>) => {
+      if (disabled || loading) return
+      const showDrip = !shadow && !ghost && effect
+      /* istanbul ignore next */
+      if (showDrip && buttonRef.current) {
+        const rect = buttonRef.current.getBoundingClientRect()
+        setDripShow(true)
+        setDripX(event.clientX - rect.left)
+        setDripY(event.clientY - rect.top)
+      }
 
-  const childrenWithIcon = useMemo(
-    () =>
-      getButtonChildrenWithIcon(auto, size, children, {
-        icon,
-        iconRight,
-      }),
-    [auto, size, children, icon, iconRight],
-  )
+      onClick && onClick(event)
+    }
 
-  return (
-    <button
-      ref={buttonRef}
-      type={htmlType}
-      className={`btn ${className}`}
-      disabled={disabled}
-      onClick={clickHandler}
-      {...props}>
-      {loading ? <ButtonLoading /> : childrenWithIcon}
-      {dripShow && (
-        <ButtonDrip x={dripX} y={dripY} color={dripColor} onCompleted={dripCompletedHandle} />
-      )}
-      <style jsx>{`
-        .btn {
-          box-sizing: border-box;
-          display: inline-block;
-          padding: 0 ${padding};
-          height: ${height};
-          line-height: ${height};
-          min-width: ${minWidth};
-          width: ${width};
-          border-radius: ${theme.layout.radius};
-          font-weight: 400;
-          font-size: ${fontSize};
-          user-select: none;
-          outline: none;
-          text-transform: capitalize;
-          justify-content: center;
-          text-align: center;
-          white-space: nowrap;
-          transition: background-color 200ms ease 0ms, box-shadow 200ms ease 0ms,
-            border 200ms ease 0ms, color 200ms ease 0ms;
-          position: relative;
-          overflow: hidden;
-          color: ${color};
-          background-color: ${bg};
-          border: 1px solid ${border};
-          cursor: ${cursor};
-          pointer-events: ${events};
-          box-shadow: ${shadow ? theme.expressiveness.shadowSmall : 'none'};
-          --zeit-ui-button-padding: ${padding};
-          --zeit-ui-button-height: ${height};
-          --zeit-ui-button-color: ${color};
-        }
+    const childrenWithIcon = useMemo(
+      () =>
+        getButtonChildrenWithIcon(auto, size, children, {
+          icon,
+          iconRight,
+        }),
+      [auto, size, children, icon, iconRight],
+    )
 
-        .btn:hover {
-          color: ${hover.color};
-          --zeit-ui-button-color: ${hover.color};
-          background-color: ${hover.bg};
-          border-color: ${hover.border};
-          cursor: ${cursor};
-          pointer-events: ${events};
-          box-shadow: ${shadow ? theme.expressiveness.shadowMedium : 'none'};
-          transform: translate3d(0px, ${shadow ? '-1px' : '0px'}, 0px);
-        }
+    return (
+      <button
+        ref={buttonRef}
+        type={htmlType}
+        className={`btn ${className}`}
+        disabled={disabled}
+        onClick={clickHandler}
+        {...props}>
+        {loading ? <ButtonLoading /> : childrenWithIcon}
+        {dripShow && (
+          <ButtonDrip x={dripX} y={dripY} color={dripColor} onCompleted={dripCompletedHandle} />
+        )}
+        <style jsx>{`
+          .btn {
+            box-sizing: border-box;
+            display: inline-block;
+            padding: 0 ${padding};
+            height: ${height};
+            line-height: ${height};
+            min-width: ${minWidth};
+            width: ${width};
+            border-radius: ${theme.layout.radius};
+            font-weight: 400;
+            font-size: ${fontSize};
+            user-select: none;
+            outline: none;
+            text-transform: capitalize;
+            justify-content: center;
+            text-align: center;
+            white-space: nowrap;
+            transition: background-color 200ms ease 0ms, box-shadow 200ms ease 0ms,
+              border 200ms ease 0ms, color 200ms ease 0ms;
+            position: relative;
+            overflow: hidden;
+            color: ${color};
+            background-color: ${bg};
+            border: 1px solid ${border};
+            cursor: ${cursor};
+            pointer-events: ${events};
+            box-shadow: ${shadow ? theme.expressiveness.shadowSmall : 'none'};
+            --zeit-ui-button-padding: ${padding};
+            --zeit-ui-button-height: ${height};
+            --zeit-ui-button-color: ${color};
+          }
 
-        .btn :global(.text) {
-          position: relative;
-          z-index: 1;
-          display: inline-flex;
-          justify-content: center;
-          align-items: center;
-          text-align: center;
-          line-height: inherit;
-          top: -1px;
-        }
+          .btn:hover {
+            color: ${hover.color};
+            --zeit-ui-button-color: ${hover.color};
+            background-color: ${hover.bg};
+            border-color: ${hover.border};
+            cursor: ${cursor};
+            pointer-events: ${events};
+            box-shadow: ${shadow ? theme.expressiveness.shadowMedium : 'none'};
+            transform: translate3d(0px, ${shadow ? '-1px' : '0px'}, 0px);
+          }
 
-        .btn :global(.text p),
-        .btn :global(.text pre),
-        .btn :global(.text div) {
-          margin: 0;
-        }
-      `}</style>
-    </button>
-  )
-}
+          .btn :global(.text) {
+            position: relative;
+            z-index: 1;
+            display: inline-flex;
+            justify-content: center;
+            align-items: center;
+            text-align: center;
+            line-height: inherit;
+            top: -1px;
+          }
 
-const MemoButton = React.memo<React.PropsWithChildren<ButtonProps>>(Button)
+          .btn :global(.text p),
+          .btn :global(.text pre),
+          .btn :global(.text div) {
+            margin: 0;
+          }
+        `}</style>
+      </button>
+    )
+  },
+)
 
-export default withDefaults(MemoButton, defaultProps)
+type ButtonComponent<T, P = {}> = React.ForwardRefExoticComponent<
+  PropsWithoutRef<P> & RefAttributes<T>
+>
+type ComponentProps = Partial<typeof defaultProps> &
+  Omit<Props, keyof typeof defaultProps> &
+  NativeAttrs
+
+Button.defaultProps = defaultProps
+
+export default React.memo(Button) as ButtonComponent<HTMLButtonElement, ComponentProps>