React polymorphic (generic, overridable) component with Typescript

react-polymorphic-component.tsx

/**
 * Resources:
 * - https://www.benmvp.com/blog/forwarding-refs-polymorphic-react-component-typescript/
 * - https://react-typescript-cheatsheet.netlify.app/docs/advanced/patterns_by_usecase/#polymorphic-components-eg-with-as-props
 * - https://react-typescript-cheatsheet.netlify.app/docs/basic/getting-started/forward_and_create_ref/
 * - https://fettblog.eu/typescript-react-generic-forward-refs/
 * - https://github.com/mui/material-ui/blob/66a69ececd29cafd55ba7d52a56442771aa46b6e/packages/mui-types/index.d.ts#L94
 */

import type { ComponentPropsWithRef, ElementType } from "react"

/**
 * Remove properties `K` from `T`.
 * Distributive for union types.
 *
 * - https://www.typescriptlang.org/docs/handbook/2/conditional-types.html#distributive-conditional-types
 * - https://stackoverflow.com/questions/57103834/typescript-omit-a-property-from-all-interfaces-in-a-union-but-keep-the-union-s
 */
export type DistributiveOmit<T, K extends keyof any> = T extends any
  ? Omit<T, K>
  : never

type As<T extends ElementType = ElementType> = { as?: T }

type PolymorphicComponentProps<
  C extends ElementType,
  Props = {},
> = As<C> &
  Props &
  DistributiveOmit<ComponentPropsWithRef<C>, keyof Props | "as">

////////////// Example ///////////////
  
import { type ElementType, type ForwardedRef, forwardRef } from "react"
import Link from "next/link"

type Props = { 
  loading?: boolean
}

type ButtonProps<C extends ElementType = "button"> =
  PolymorphicComponentProps<C, Props>

export const Button = forwardRef(function Button<
  C extends ElementType = "button",
>({ as, ...rest }: ButtonProps<C>, ref: ForwardedRef<ElementType<C>>) {
  const Component = as || "button"

  return (
    <Component
      ref={ref}
      {...rest}
    >
      Click me
    </Component>
  )
}) as <C extends ElementType = "button">(
  props: ButtonProps<C>,
) => JSX.Element | null

const Component = () => (
  <>
    <Button as="span" />
    <Button as={Link} href="https://google.com" />
  </>
)