Support inference extra props from as props

Author: gtkatakura

package.json

@@ -51,6 +51,7 @@
     "react-dom": "^16.13.1",
     "react-scripts": "^3.4.1",
     "reakit": "^1.1.2",
+    "reakit-utils": "^0.13.1",
     "theme-ui": "^0.3.1",
     "ts-loader": "^8.0.1",
     "tsdx": "^0.13.2",

src/Button.stories.tsx

@@ -1,8 +1,7 @@
-import React from 'react'
+import React, { useReducer, useRef, useEffect, forwardRef, Ref } from 'react'
 import { ThemeProvider } from 'theme-ui'
 import { withA11y } from '@storybook/addon-a11y'
 import { withKnobs, text, color, boolean } from '@storybook/addon-knobs'
-import { Checkbox, useCheckboxState } from 'reakit'
 
 import Button from '.'
 
@@ -46,8 +45,55 @@ export function TheSXProp() {
   )
 }
 
-export function TheAsProp() {
-  const checkbox = useCheckboxState()
+export function TheAsPropWithIntrinsicElement() {
+  const ref = useRef<HTMLAnchorElement>(null)
+  const [checked, toggle] = useReducer((value) => !value, false)
+
+  useEffect(() => {
+    console.warn(
+      `TheAsPropWithIntrinsicElement: Detects properly type of \`ref\` from \`as\``,
+      ref.current
+    )
+  }, [])
+
+  return (
+    <ThemeProvider
+      theme={{
+        colors: {
+          background: '#FFFFFF',
+          primary: '#2F323A',
+          secondary: '#4F5D75',
+        },
+      }}
+    >
+      <Button ref={ref} as="a" target="_blank" onClick={toggle}>
+        {checked ? '😄 Happy' : '😞 Sad'}
+      </Button>
+    </ThemeProvider>
+  )
+}
+
+const Link = forwardRef(function Link(
+  { children, ...props }: JSX.IntrinsicElements['a'],
+  ref: Ref<HTMLAnchorElement>
+) {
+  return (
+    <a ref={ref} {...props}>
+      {children}
+    </a>
+  )
+})
+
+export function TheAsPropWithCustomElement() {
+  const ref = useRef<HTMLAnchorElement>(null)
+  const [checked, toggle] = useReducer((value) => !value, false)
+
+  useEffect(() => {
+    console.warn(
+      `TheAsPropWithCustomElement: Detects properly type of \`ref\` from \`as\``,
+      ref.current
+    )
+  }, [])
 
   return (
     <ThemeProvider
@@ -59,9 +105,9 @@ export function TheAsProp() {
         },
       }}
     >
-      <Checkbox {...checkbox} as={Button}>
-        {checkbox.state ? '😄 Happy' : '😞 Sad'}
-      </Checkbox>
+      <Button ref={ref} as={Link} target="_blank" onClick={toggle}>
+        {checked ? '😄 Happy' : '😞 Sad'}
+      </Button>
     </ThemeProvider>
   )
 }

src/index.tsx

@@ -1,15 +1,29 @@
-import React, { forwardRef, Ref } from 'react'
+import React, { forwardRef, Ref, PropsWithoutRef } from 'react'
 import { Button as A11yButton, ButtonProps as A11yProps } from 'reakit'
+import { PropsWithAs, As } from 'reakit-utils/types'
 import {
   Button as ThemeAwareButton,
   ButtonProps as ThemeAwareProps,
 } from 'theme-ui'
 
-type Props = A11yProps & ThemeAwareProps
+type Component<DefaultAs extends As, DefaultProps> = {
+  <T extends As = DefaultAs>(
+    props: PropsWithAs<PropsWithoutRef<DefaultProps>, T>
+  ): JSX.Element
+}
+
+type ButtonProps = A11yProps & { as?: As }
 
-function Button({ as, ...props }: Props, ref: Ref<HTMLButtonElement>) {
-  return <A11yButton {...props} ref={ref} as={as ?? ThemeAwareButton} />
+const Button = (
+  { as = ThemeAwareButton, ...props }: ButtonProps,
+  ref: Ref<As>
+) => {
+  return <A11yButton ref={ref} as={as} {...props} />
 }
 
-export { A11yProps, ThemeAwareProps }
-export default forwardRef(Button)
+export { ButtonProps, A11yProps, ThemeAwareProps }
+
+export default forwardRef(Button) as Component<
+  typeof ThemeAwareButton,
+  A11yProps
+>