Support inference extra props from as props

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

@@ -4,12 +4,17 @@ import {
   Button as ThemeAwareButton,
   ButtonProps as ThemeAwareProps,
 } from 'theme-ui'
+import { As, Props, Component } from '@vtex-components/types'
 
-type Props = A11yProps & ThemeAwareProps
+type ButtonProps<T extends As = typeof ThemeAwareButton> = Props<T, A11yProps>
 
-function Button({ as, ...props }: Props, ref: Ref<HTMLButtonElement>) {
-  return <A11yButton {...props} ref={ref} as={as ?? ThemeAwareButton} />
+const Button = (
+  { as = ThemeAwareButton, ...props }: ButtonProps<As>,
+  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<ButtonProps>

src/typings.d.ts

@@ -15,3 +15,27 @@ declare module '*.svg' {
   export default svgUrl
   export { svgComponent as ReactComponent }
 }
+
+declare module '@vtex-components/types' {
+  import { PropsWithoutRef } from 'react'
+  import { As, PropsWithAs } from 'reakit-utils/types'
+
+  export { As }
+
+  export type Props<T extends As, BaseProps> = PropsWithAs<
+    PropsWithoutRef<BaseProps>,
+    T
+  >
+
+  export type AsOf<T> = T extends Props<infer TAs, any> ? TAs : never
+
+  export type BasePropsOf<T> = T extends Props<any, infer TProps>
+    ? TProps
+    : never
+
+  export type Component<DefaultProps extends Props<any, any>> = {
+    <T extends As = AsOf<DefaultProps>>(
+      props: Props<T, BasePropsOf<DefaultProps>>
+    ): JSX.Element
+  }
+}