React Polymorphic Component

Polymorphic.tsx

import React from 'react';

export type ExtendHTMLProps<P extends object = {}> = Merge<React.HTMLAttributes<HTMLElement>, P>;
export type Merge<A, B> = Omit<A, keyof B> & B;

type ExtendAs<C extends React.ElementType = 'div'> = {
  as?: C;
};

export type PolymorphicExtendedProps<
  C extends React.ElementType,
  Props = {},
> = ExtendHTMLProps<Merge<Props, ExtendAs<C>>> &
  Omit<React.ComponentPropsWithoutRef<C>, keyof Props> & {
    ref?: React.ComponentPropsWithRef<C>['ref'];
  };

export interface PolymorphicComponent<
  P extends ExtendHTMLProps = {},
  T extends React.ElementType = 'div',
> {
  <C extends React.ElementType = T>(
    props: Merge<P, PolymorphicExtendedProps<C, ExtendAs<C>>>,
  ): JSX.Element | null;

  displayName?: string;
}

export type PolymorphicProps<
  C extends React.ElementType = 'div',
  P extends ExtendAs<C> = {},
> = PolymorphicExtendedProps<C, P>;

const Polymorphic = React.forwardRef(
  <C extends React.ElementType = 'div'>(
    props: PolymorphicProps<C>,
    ref: React.ForwardedRef<HTMLElement>,
  ) => {
    const { as: Component = 'div', ...rest } = props;
    return <Component {...rest} ref={ref} />;
  },
) as PolymorphicComponent;

Polymorphic.displayName = 'Polymorphic';

export default Polymorphic;