pickVariantStyle.ts

type ThemedCss = FlattenInterpolation<ThemeProps<DefaultTheme>>;

type VariantOptions<T extends string> = {
  variant: T | undefined;
  variants: Record<T, ThemedCss>;
  defaultStyle: ThemedCss;
};

type VariantCallback<T, K extends string> = (props: T) => VariantOptions<K>;

function pickVariantStyleV2<T, K extends string>(
  fn: VariantCallback<T, K>
): (props: T) => ThemedCss {
  return function themedVariant(props: T): ThemedCss {
    const { variant, variants, defaultStyle } = fn(props);

    if (!variant) {
      return defaultStyle;
    }

    return variants[variant] || defaultStyle;
  };
}

Example usage:

const Medium = css`
  font-size: 24px;
`;

const Small = css`
  font-size: 16px;
`;

const ExtraSmall = css`
  font-size: 8px;
`;

const AwesomeComponent = styled.div`
  ${pickVariantStyle<AwesomeComponentProps, VariantsType>((props: AwesomeComponentProps) => ({
    variant: props?.size,
    variants: {
      m: Medium,
      s: Small,
      xs: ExtraSmall,
    },
    defaultStyle: Medium,
  }))}
`;