Avoid tokens weird type checking/eliding

.changeset/gentle-badgers-confess.md

@@ -0,0 +1,7 @@
+---
+'@primer/react': patch
+---
+
+Token: Update component type to be PolymorphicForwardRefComponent.
+
+this avoids types being swallowed by forwardRef (which isn't polymorphic)

src/Token/AvatarToken.tsx

@@ -4,6 +4,7 @@ import {get} from '../constants'
 import {TokenBaseProps, defaultTokenSize, tokenSizes, TokenSizeKeys} from './TokenBase'
 import Token from './Token'
 import Avatar from '../Avatar'
+import {ForwardRefComponent as PolymorphicForwardRefComponent} from '../utils/polymorphic'
 
 // TODO: update props to only accept 'large' and 'xlarge' on the next breaking change
 export interface AvatarTokenProps extends TokenBaseProps {
@@ -20,7 +21,7 @@ const AvatarContainer = styled.span<{avatarSize: TokenSizeKeys}>`
   width: ${props => `calc(${tokenSizes[props.avatarSize]} - var(--spacing))`};
 `
 
-const AvatarToken = forwardRef<HTMLElement, AvatarTokenProps>(({avatarSrc, id, size, ...rest}, forwardedRef) => {
+const AvatarToken = forwardRef(({avatarSrc, id, size, ...rest}, forwardedRef) => {
   return (
     <Token
       leadingVisual={() => (
@@ -44,7 +45,7 @@ const AvatarToken = forwardRef<HTMLElement, AvatarTokenProps>(({avatarSrc, id, s
       ref={forwardedRef}
     />
   )
-})
+}) as PolymorphicForwardRefComponent<'span' | 'a' | 'button', AvatarTokenProps>
 
 AvatarToken.defaultProps = {
   size: defaultTokenSize,

src/Token/IssueLabelToken.tsx

@@ -5,6 +5,7 @@ import RemoveTokenButton from './_RemoveTokenButton'
 import {parseToHsla, parseToRgba} from 'color2k'
 import {useTheme} from '../ThemeProvider'
 import TokenTextContainer from './_TokenTextContainer'
+import {ForwardRefComponent as PolymorphicForwardRefComponent} from '../utils/polymorphic'
 
 export interface IssueLabelTokenProps extends TokenBaseProps {
   /**
@@ -39,7 +40,7 @@ const darkModeStyles = {
     'hsla(var(--label-h), calc(var(--label-s) * 1%),calc((var(--label-l) + var(--lighten-by)) * 1%),var(--border-alpha))',
 }
 
-const IssueLabelToken = forwardRef<HTMLElement, IssueLabelTokenProps>((props, forwardedRef) => {
+const IssueLabelToken = forwardRef((props, forwardedRef) => {
   const {as, fillColor = '#999', onRemove, id, isSelected, text, size, hideRemoveButton, href, onClick, ...rest} = props
   const interactiveTokenProps = {
     as,
@@ -134,7 +135,7 @@ const IssueLabelToken = forwardRef<HTMLElement, IssueLabelTokenProps>((props, fo
       ) : null}
     </TokenBase>
   )
-})
+}) as PolymorphicForwardRefComponent<'span' | 'a' | 'button', IssueLabelTokenProps>
 
 IssueLabelToken.defaultProps = {
   fillColor: '#999',

src/Token/Token.tsx

@@ -1,14 +1,15 @@
 import React, {forwardRef, MouseEventHandler} from 'react'
 import Box from '../Box'
-import {merge, SxProp} from '../sx'
+import {BetterSystemStyleObject, merge, SxProp} from '../sx'
 import {defaultSxProp} from '../utils/defaultSxProp'
 import TokenBase, {defaultTokenSize, isTokenInteractive, TokenBaseProps} from './TokenBase'
 import RemoveTokenButton from './_RemoveTokenButton'
 import TokenTextContainer from './_TokenTextContainer'
+import {ForwardRefComponent as PolymorphicForwardRefComponent} from '../utils/polymorphic'
 
 // Omitting onResize and onResizeCapture because seems like React 18 types includes these menthod in the expansion but React 17 doesn't.
 // TODO: This is a temporary solution until we figure out why these methods are causing type errors.
-export interface TokenProps extends Omit<TokenBaseProps, 'onResize' | 'onResizeCapture'> {
+export interface TokenProps extends TokenBaseProps, SxProp {
   /**
    * A function that renders a component before the token text
    */
@@ -30,83 +31,81 @@ const LeadingVisualContainer: React.FC<React.PropsWithChildren<Pick<TokenBasePro
   </Box>
 )
 
-const Token = forwardRef<HTMLAnchorElement | HTMLButtonElement | HTMLSpanElement, TokenProps & SxProp>(
-  (props, forwardedRef) => {
-    const {
-      as,
-      onRemove,
-      id,
-      leadingVisual: LeadingVisual,
-      text,
-      size,
-      hideRemoveButton,
-      href,
-      onClick,
-      sx: sxProp = defaultSxProp,
-      ...rest
-    } = props
-    const hasMultipleActionTargets = isTokenInteractive(props) && Boolean(onRemove) && !hideRemoveButton
-    const onRemoveClick: MouseEventHandler = e => {
-      e.stopPropagation()
-      onRemove && onRemove()
-    }
-    const interactiveTokenProps = {
-      as,
-      href,
-      onClick,
-    }
-    const sx = merge(
-      {
-        backgroundColor: 'neutral.subtle',
-        borderColor: props.isSelected ? 'fg.default' : 'border.subtle',
-        borderStyle: 'solid',
-        borderWidth: `${tokenBorderWidthPx}px`,
-        color: props.isSelected ? 'fg.default' : 'fg.muted',
-        maxWidth: '100%',
-        paddingRight: !(hideRemoveButton || !onRemove) ? 0 : undefined,
-        ...(isTokenInteractive(props)
-          ? {
-              '&:hover': {
-                backgroundColor: 'neutral.muted',
-                boxShadow: 'shadow.medium',
-                color: 'fg.default',
-              },
-            }
-          : {}),
-      },
-      sxProp as SxProp,
-    )
+const Token = forwardRef((props, forwardedRef) => {
+  const {
+    as,
+    onRemove,
+    id,
+    leadingVisual: LeadingVisual,
+    text,
+    size,
+    hideRemoveButton,
+    href,
+    onClick,
+    sx: sxProp = defaultSxProp,
+    ...rest
+  } = props
+  const hasMultipleActionTargets = isTokenInteractive(props) && Boolean(onRemove) && !hideRemoveButton
+  const onRemoveClick: MouseEventHandler = e => {
+    e.stopPropagation()
+    onRemove && onRemove()
+  }
+  const interactiveTokenProps = {
+    as,
+    href,
+    onClick,
+  }
+  const sx = merge<BetterSystemStyleObject>(
+    {
+      backgroundColor: 'neutral.subtle',
+      borderColor: props.isSelected ? 'fg.default' : 'border.subtle',
+      borderStyle: 'solid',
+      borderWidth: `${tokenBorderWidthPx}px`,
+      color: props.isSelected ? 'fg.default' : 'fg.muted',
+      maxWidth: '100%',
+      paddingRight: !(hideRemoveButton || !onRemove) ? 0 : undefined,
+      ...(isTokenInteractive(props)
+        ? {
+            '&:hover': {
+              backgroundColor: 'neutral.muted',
+              boxShadow: 'shadow.medium',
+              color: 'fg.default',
+            },
+          }
+        : {}),
+    },
+    sxProp,
+  )
 
-    return (
-      <TokenBase
-        onRemove={onRemove}
-        id={id?.toString()}
-        text={text}
-        size={size}
-        sx={sx}
-        {...(!hasMultipleActionTargets ? interactiveTokenProps : {})}
-        {...rest}
-        ref={forwardedRef}
-      >
-        {LeadingVisual ? (
-          <LeadingVisualContainer size={size}>
-            <LeadingVisual />
-          </LeadingVisualContainer>
-        ) : null}
-        <TokenTextContainer {...(hasMultipleActionTargets ? interactiveTokenProps : {})}>{text}</TokenTextContainer>
-        {!hideRemoveButton && onRemove ? (
-          <RemoveTokenButton
-            borderOffset={tokenBorderWidthPx}
-            onClick={onRemoveClick}
-            size={size}
-            isParentInteractive={isTokenInteractive(props)}
-            aria-hidden={hasMultipleActionTargets ? 'true' : 'false'}
-          />
-        ) : null}
-      </TokenBase>
-    )
-  },
-)
+  return (
+    <TokenBase
+      onRemove={onRemove}
+      id={id?.toString()}
+      text={text}
+      size={size}
+      sx={sx}
+      {...(!hasMultipleActionTargets ? interactiveTokenProps : {})}
+      {...rest}
+      ref={forwardedRef}
+    >
+      {LeadingVisual ? (
+        <LeadingVisualContainer size={size}>
+          <LeadingVisual />
+        </LeadingVisualContainer>
+      ) : null}
+      <TokenTextContainer {...(hasMultipleActionTargets ? interactiveTokenProps : {})}>{text}</TokenTextContainer>
+      {!hideRemoveButton && onRemove ? (
+        <RemoveTokenButton
+          borderOffset={tokenBorderWidthPx}
+          onClick={onRemoveClick}
+          size={size}
+          isParentInteractive={isTokenInteractive(props)}
+          aria-hidden={hasMultipleActionTargets ? 'true' : 'false'}
+        />
+      ) : null}
+    </TokenBase>
+  )
+}) as PolymorphicForwardRefComponent<'a' | 'button' | 'span', TokenProps>
 
 Token.displayName = 'Token'
 

src/Token/TokenBase.tsx

@@ -1,8 +1,9 @@
-import React, {KeyboardEvent} from 'react'
+import React, {ComponentProps, KeyboardEvent} from 'react'
 import styled from 'styled-components'
 import {variant} from 'styled-system'
 import {get} from '../constants'
 import sx, {SxProp} from '../sx'
+import {ForwardRefComponent as PolymorphicForwardRefComponent} from '../utils/polymorphic'
 
 // TODO: remove invalid "extralarge" size name in next breaking change
 /** @deprecated 'extralarge' to be removed to align with size naming ADR https://github.com/github/primer/blob/main/adrs/2022-02-09-size-naming-guidelines.md **/
@@ -50,14 +51,12 @@ export interface TokenBaseProps
   size?: TokenSizeKeys
 }
 
-type TokenElements = HTMLSpanElement | HTMLButtonElement | HTMLAnchorElement
-
 export const isTokenInteractive = ({
   as = 'span',
   onClick,
   onFocus,
   tabIndex = -1,
-}: Pick<TokenBaseProps, 'as' | 'onClick' | 'onFocus' | 'tabIndex'>) =>
+}: Pick<ComponentProps<typeof TokenBase>, 'as' | 'onClick' | 'onFocus' | 'tabIndex'>) =>
   Boolean(onFocus || onClick || tabIndex > -1 || ['a', 'button'].includes(as))
 
 const xlargeVariantStyles = {
@@ -133,25 +132,23 @@ const StyledTokenBase = styled.span<SxProp>`
   ${sx}
 `
 
-const TokenBase = React.forwardRef<TokenElements, TokenBaseProps & SxProp>(
-  ({text, onRemove, onKeyDown, id, ...rest}, forwardedRef) => {
-    return (
-      <StyledTokenBase
-        onKeyDown={(event: KeyboardEvent<TokenElements>) => {
-          onKeyDown && onKeyDown(event)
+const TokenBase = React.forwardRef(({text, onRemove, onKeyDown, id, ...rest}, forwardedRef) => {
+  return (
+    <StyledTokenBase
+      onKeyDown={(event: KeyboardEvent<HTMLSpanElement & HTMLAnchorElement & HTMLButtonElement>) => {
+        onKeyDown && onKeyDown(event)
 
-          if ((event.key === 'Backspace' || event.key === 'Delete') && onRemove) {
-            onRemove()
-          }
-        }}
-        aria-label={onRemove ? `${text}, press backspace or delete to remove` : undefined}
-        id={id?.toString()}
-        {...rest}
-        ref={forwardedRef}
-      />
-    )
-  },
-)
+        if ((event.key === 'Backspace' || event.key === 'Delete') && onRemove) {
+          onRemove()
+        }
+      }}
+      aria-label={onRemove ? `${text}, press backspace or delete to remove` : undefined}
+      id={id?.toString()}
+      {...rest}
+      ref={forwardedRef}
+    />
+  )
+}) as PolymorphicForwardRefComponent<'span' | 'a' | 'button', TokenBaseProps & SxProp>
 
 TokenBase.defaultProps = {
   as: 'span',