TypeScript implementation of a styled-components-like API for material-ui's withStyles

mui-styled.tsx

import React from 'react'
import {
  withStyles,
  Theme,
  StyledComponentProps,
} from '@material-ui/core/styles'
import cx from 'classnames'
import {
  WithStylesOptions,
  CSSProperties,
} from '@material-ui/core/styles/withStyles'

type ComponentProps<
  T extends React.ComponentType<any> | keyof JSX.IntrinsicElements
> = T extends React.ComponentType<infer P>
  ? P
  : T extends keyof JSX.IntrinsicElements ? JSX.IntrinsicElements[T] : {}

type RefTypeOfIntrinsic<T> = T extends { ref?: infer R }
  ? {} extends R ? never : R
  : never

type RefTypeOfComponent<
  T extends React.ComponentType<any> | keyof JSX.IntrinsicElements
> = T extends (props: infer P) => unknown
  ? 'ref' extends keyof P ? (P extends { ref?: infer R } ? R : never) : never
  : T extends new (props: never) => infer I
    ? I
    : T extends keyof JSX.IntrinsicElements
      ? RefTypeOfIntrinsic<JSX.IntrinsicElements[T]>
      : never

interface StyledProps extends Pick<StyledComponentProps<'root'>, 'classes'> {
  className?: string
}

export default function styled<
  T extends React.ComponentType<any> | keyof JSX.IntrinsicElements
>(
  Component: 'className' extends keyof ComponentProps<T>
    ? ComponentProps<T> extends { className?: string } ? T : never
    : never,
  innerRef?: boolean
) {
  const refKeyName = innerRef ? 'innerRef' : 'ref'
  return styled
  function styled(
    style: ((theme: Theme) => CSSProperties) | CSSProperties,
    options?: WithStylesOptions<'root'>
  ) {
    Styled.displayName = getComponentName(Component)
    const StyledComponent = React.forwardRef(Styled)

    const styles =
      typeof style === 'function'
        ? (theme: Theme) => ({ root: style(theme) })
        : { root: style }

    const result = withStyles(styles, {
      ...options,
      withTheme: false,
      name: Styled.displayName,
    })((StyledComponent as unknown) as React.SFC<StyledProps>)

    result.displayName = Styled.displayName
      ? `styled(${Styled.displayName})`
      : `styled`

    return (result as unknown) as React.ComponentClass<
      Pick<
        ComponentProps<typeof Component>,
        Exclude<keyof ComponentProps<typeof Component>, 'innerRef' | 'ref'>
      > &
        (never extends RefTypeOfComponent<typeof Component>
          ? {}
          : { innerRef?: React.Ref<RefTypeOfComponent<typeof Component>> }) & {
          ref?: never
        }
    >

    function Styled(
      props: ComponentProps<T> & StyledProps,
      ref?: React.Ref<any>
    ) {
      const { classes, className, ...other } = props as StyledProps
      return React.createElement(Component, {
        className: cx(classes!.root, className),
        ...other,
        [refKeyName]: ref,
      })
    }
  }
}

function getComponentName(Component: string | React.ComponentType<any>) {
  if (typeof Component === 'string') {
    return Component
  }
  const name = Component.displayName || Component.name
  const m = /^WithStyles\((.*)\)$/.exec(name)
  return m !== null ? m[1] : name
}