feat: polymoprphic theme component

src/Theme/Theme.stories.tsx

@@ -12,7 +12,9 @@ export default {
   component: Theme,
 } as Meta
 
-export const Default: Story<ThemeProps> = (args) => {
+export const Default: Story<ThemeProps<React.ElementType<HTMLDivElement>>> = (
+  args
+) => {
   const { theme, setTheme } = useTheme()
 
   return (
@@ -43,7 +45,24 @@ export const Default: Story<ThemeProps> = (args) => {
 }
 Default.args = {}
 
-export const NestedThemes: Story<ThemeProps> = (args) => {
+export const AsCustomTag: Story<
+  ThemeProps<React.ElementType<HTMLBodyElement>>
+> = (args) => {
+  const { theme, setTheme } = useTheme()
+
+  return (
+    <Theme dataTheme={theme} as="body">
+      <div>
+        <ThemeItem dataTheme={theme} />
+      </div>
+    </Theme>
+  )
+}
+AsCustomTag.args = {}
+
+export const NestedThemes: Story<
+  ThemeProps<React.ElementType<HTMLDivElement>>
+> = (args) => {
   const { theme, setTheme } = useTheme()
 
   const renderNestedThemes = (themes: readonly string[]) => {

src/Theme/Theme.tsx

@@ -1,27 +1,90 @@
-import React, { MutableRefObject, useEffect, useRef, useState } from 'react'
+import React, {
+  MutableRefObject,
+  ElementType,
+  useEffect,
+  useRef,
+  useState,
+  ComponentPropsWithoutRef,
+} from 'react'
 import { defaultTheme } from '../constants'
 
 import { DataTheme, IComponentBaseProps } from '../types'
 import { ThemeContext } from './ThemeContext'
 import { getThemeFromClosestAncestor } from './utils'
 
-export type ThemeProps = Omit<
-  React.HTMLAttributes<HTMLDivElement>,
-  'onChange'
-> &
-  IComponentBaseProps & {
-    onChange?: (theme: DataTheme) => void
-  }
+// Polymorphic component with forwardable refs types from: https://www.benmvp.com/blog/forwarding-refs-polymorphic-react-component-typescript/
+
+// Source: https://github.com/emotion-js/emotion/blob/master/packages/styled-base/types/helper.d.ts
+// A more precise version of just React.ComponentPropsWithoutRef on its own
+export type PropsOf<
+  C extends keyof JSX.IntrinsicElements | React.JSXElementConstructor<any>
+> = JSX.LibraryManagedAttributes<C, React.ComponentPropsWithoutRef<C>>
+
+type AsProp<C extends React.ElementType> = {
+  /**
+   * An override of the default HTML tag.
+   * Can also be another React component.
+   */
+  as?: C
+}
+
+/**
+ * Allows for extending a set of props (`ExtendedProps`) by an overriding set of props
+ * (`OverrideProps`), ensuring that any duplicates are overridden by the overriding
+ * set of props.
+ */
+type ExtendableProps<ExtendedProps = {}, OverrideProps = {}> = OverrideProps &
+  Omit<ExtendedProps, keyof OverrideProps>
+
+/**
+ * Allows for inheriting the props from the specified element type so that
+ * props like children, className & style work, as well as element-specific
+ * attributes like aria roles. The component (`C`) must be passed in.
+ */
+type InheritableElementProps<
+  C extends React.ElementType,
+  Props = {}
+> = ExtendableProps<PropsOf<C>, Props>
+
+/**
+ * A more sophisticated version of `InheritableElementProps` where
+ * the passed in `as` prop will determine which props can be included
+ */
+type PolymorphicComponentProps<
+  C extends React.ElementType,
+  Props = {}
+> = InheritableElementProps<C, Props & AsProp<C>>
+
+/** * Utility type to extract the `ref` prop from a polymorphic component */
+type PolymorphicRef<C extends React.ElementType> =
+  React.ComponentPropsWithRef<C>['ref']
+/** * A wrapper of `PolymorphicComponentProps` that also includes the `ref` * prop for the polymorphic component */
+type PolymorphicComponentPropsWithRef<
+  C extends React.ElementType,
+  Props = {}
+> = PolymorphicComponentProps<C, Props> & { ref?: PolymorphicRef<C> }
+
+interface Props {
+  children?: React.ReactNode
+  onChange?: (theme: DataTheme) => void
+}
+
+export type ThemeProps<C extends React.ElementType> =
+  PolymorphicComponentPropsWithRef<C, Props> & IComponentBaseProps
+
+type ThemeComponent = <C extends React.ElementType = 'div'>(
+  props: ThemeProps<C>
+) => React.ReactElement | null
+
+const Theme: ThemeComponent = React.forwardRef(
+  <C extends React.ElementType = 'div'>(
+    { as, children, dataTheme, onChange, ...props }: ThemeProps<C>,
+    ref?: PolymorphicRef<C>
+  ) => {
+    const Component = as || 'div'
 
-const Theme = React.forwardRef<HTMLDivElement, ThemeProps>(
-  (
-    { children, dataTheme, onChange, className, ...props },
-    ref
-  ): JSX.Element => {
     // Either use provided ref or create a new ref
-    const themeRef = useRef<HTMLDivElement>(
-      (ref as MutableRefObject<HTMLDivElement>)?.current
-    )
+    const themeRef = useRef<PolymorphicRef<C>>(ref?.current)
 
     const closestAncestorTheme = getThemeFromClosestAncestor(themeRef)
 
@@ -46,11 +109,12 @@ const Theme = React.forwardRef<HTMLDivElement, ThemeProps>(
 
     return (
       <ThemeContext.Provider value={{ theme, setTheme: handleThemeChange }}>
-        <div {...props} data-theme={theme} className={className} ref={themeRef}>
+        <Component {...props} data-theme={theme} ref={themeRef}>
           {children}
-        </div>
+        </Component>
       </ThemeContext.Provider>
     )
   }
 )
+
 export default Theme

src/Theme/index.tsx

@@ -1,3 +1,3 @@
 import Theme, { ThemeProps as TThemeProps } from './Theme'
-export type ThemeProps = TThemeProps
+export type ThemeProps<T extends React.ElementType> = TThemeProps<T>
 export default Theme

src/index.ts

@@ -245,5 +245,5 @@ export type WindowMockupProps = TWindowMockupProps
 // Utils > Theme
 export { default as Theme } from './Theme'
 import { ThemeProps as TThemeProps } from './Theme'
-export type ThemeProps = TThemeProps
+export type ThemeProps<T extends React.ElementType> = TThemeProps<T>
 export { useTheme } from './Theme/useTheme'