feat(v3): support asChild composition (#8302)

* feat: support as child composition

* refactor: simplify types

* chore: deprecate as props

* refactor: rename vars

.changeset/ninety-humans-call.md

@@ -0,0 +1,5 @@
+---
+"@chakra-ui/react": minor
+---
+
+Add support for `asChild` in chakra factory

MIGRATION.md

@@ -0,0 +1,23 @@
+# v3 Migration Guide
+
+## Changed
+
+## Added
+
+### `asChild` prop
+
+Removed support for `as` prop due to the type complexity involved.
+
+**Action:** Replace `asChild` in `chakra` factory and existing components.
+
+```tsx
+import { Button } from "@chakra-ui/react"
+
+const Demo = () => {
+  return (
+    <Button asChild>
+      <a href="#">Child</a>
+    </Button>
+  )
+}
+```

packages/components/src/system/factory.ts

@@ -1,9 +1,10 @@
+import { ElementType } from "react"
 import { ChakraStyledOptions, HTMLChakraComponents, styled } from "./system"
-import { As, ChakraComponent } from "./system.types"
+import { ChakraComponent } from "./system.types"
 import { DOMElements } from "./system.utils"
 
-type ChakraFactory = {
-  <T extends As, P extends object = {}>(
+interface ChakraFactory {
+  <T extends ElementType, P extends object = {}>(
     component: T,
     options?: ChakraStyledOptions,
   ): ChakraComponent<T, P>

packages/components/src/system/forward-ref.tsx

@@ -2,14 +2,14 @@
  * All credit goes to Chance (Reach UI), Haz (Reakit) and (fluentui)
  * for creating the base type definitions upon which we improved on
  */
-import { forwardRef as forwardReactRef } from "react"
-import { As, ComponentWithAs, PropsOf, RightJoinProps } from "./system.types"
+import { ElementType, forwardRef as forwardReactRef } from "react"
+import { ComponentWithAs, PropsOf, RightJoinProps } from "./system.types"
 
-export function forwardRef<Props extends object, Component extends As>(
+export function forwardRef<Props extends object, Component extends ElementType>(
   component: React.ForwardRefRenderFunction<
     any,
     RightJoinProps<PropsOf<Component>, Props> & {
-      as?: As
+      as?: ElementType
     }
   >,
 ) {

packages/components/src/system/merge-props.ts

@@ -0,0 +1,61 @@
+import { callAll } from "@chakra-ui/utils"
+
+interface Props {
+  [key: string]: any
+}
+
+const clsx = (...args: (string | undefined)[]) =>
+  args
+    .map((str) => str?.trim?.())
+    .filter(Boolean)
+    .join(" ")
+
+type TupleTypes<T extends any[]> = T[number]
+
+type UnionToIntersection<U> = (U extends any ? (k: U) => void : never) extends (
+  k: infer I,
+) => void
+  ? I
+  : never
+
+const eventRegex = /^on[A-Z]/
+
+export function mergeProps<T extends Props>(
+  ...args: T[]
+): UnionToIntersection<TupleTypes<T[]>> {
+  let result: Props = {}
+
+  for (let props of args) {
+    for (let key in result) {
+      if (
+        eventRegex.test(key) &&
+        typeof result[key] === "function" &&
+        typeof props[key] === "function"
+      ) {
+        result[key] = callAll(result[key], props[key])
+        continue
+      }
+
+      if (key === "className" || key === "class") {
+        result[key] = clsx(result[key], props[key])
+        continue
+      }
+
+      if (key === "style") {
+        result[key] = Object.assign({}, result[key] ?? {}, props[key] ?? {})
+        continue
+      }
+
+      result[key] = props[key] !== undefined ? props[key] : result[key]
+    }
+
+    // Add props from b that are not in a
+    for (let key in props) {
+      if (result[key] === undefined) {
+        result[key] = props[key]
+      }
+    }
+  }
+
+  return result as any
+}

packages/components/src/system/system.stories.tsx

@@ -8,6 +8,14 @@ export default {
 
 const MotionBox = motion(chakra.div)
 
+export const WithAsChild = () => {
+  return (
+    <chakra.button bg="red.200" asChild>
+      <a href="dfd">sdfsd</a>
+    </chakra.button>
+  )
+}
+
 export const WithFramerMotion = () => (
   <MotionBox
     mt="40px"

packages/components/src/system/system.ts

@@ -1,3 +1,4 @@
+import { mergeRefs } from "@chakra-ui/hooks/use-merge-refs"
 import {
   css,
   isStylePropFn,
@@ -11,10 +12,24 @@ import { runIfFn } from "@chakra-ui/utils/run-if-fn"
 import { splitProps } from "@chakra-ui/utils/split-props"
 import { Dict } from "@chakra-ui/utils/types"
 import createStyled, { CSSObject, FunctionInterpolation } from "@emotion/styled"
-import { createElement, forwardRef } from "react"
+import {
+  Children,
+  createElement,
+  ElementType,
+  forwardRef,
+  isValidElement,
+  useMemo,
+} from "react"
 import { useColorMode } from "../color-mode"
+import { mergeProps } from "./merge-props"
 import { shouldForwardProp } from "./should-forward-prop"
-import { As, ChakraComponent, ChakraProps, PropsOf } from "./system.types"
+import {
+  AsChildProps,
+  AsProps,
+  ChakraComponent,
+  ChakraProps,
+  PropsOf,
+} from "./system.types"
 import { DOMElements } from "./system.utils"
 
 const emotion_styled = interopDefault(createStyled)
@@ -74,7 +89,7 @@ export interface ChakraStyledOptions extends Dict {
     | ((props: StyleResolverProps) => SystemStyleObject)
 }
 
-export function styled<T extends As, P extends object = {}>(
+export function styled<T extends ElementType, P extends object = {}>(
   component: T,
   options?: ChakraStyledOptions,
 ) {
@@ -85,19 +100,58 @@ export function styled<T extends As, P extends object = {}>(
   }
 
   const styleObject = toCSSObject({ baseStyle })
+
   const Component = emotion_styled(
     component as React.ComponentType<any>,
     styledOptions,
   )(styleObject)
 
-  const chakraComponent = forwardRef(function ChakraComponent(props, ref) {
-    const { colorMode, forced } = useColorMode()
-    return createElement(Component, {
-      ref,
-      "data-theme": forced ? colorMode : undefined,
-      ...props,
-    })
-  })
+  const chakraComponent = forwardRef<any, any>(
+    function ChakraComponent(props, ref) {
+      const { asChild, children, ...restProps } = props
+
+      const { colorMode, forced } = useColorMode()
+
+      const dataTheme = forced ? colorMode : undefined
+
+      if (!asChild) {
+        return createElement(
+          Component,
+          {
+            ref,
+            "data-theme": dataTheme,
+            ...restProps,
+          },
+          children,
+        )
+      }
+
+      const onlyChild = Children.only(props.children)
+
+      if (isValidElement(onlyChild)) {
+        const composedProps = mergeProps(restProps, onlyChild.props ?? {})
+
+        const composedRef = ref
+          ? mergeRefs(ref, (onlyChild as any).ref)
+          : (onlyChild as any).ref
+
+        // eslint-disable-next-line react-hooks/rules-of-hooks
+        const styledElement = useMemo(
+          () =>
+            emotion_styled(onlyChild.type as any, styledOptions)(styleObject),
+          [onlyChild.type],
+        )
+
+        return createElement(styledElement, {
+          ref: composedRef,
+          "data-theme": dataTheme,
+          ...composedProps,
+        })
+      }
+
+      return onlyChild
+    },
+  )
 
   return chakraComponent as ChakraComponent<T, P>
 }
@@ -106,8 +160,10 @@ export type HTMLChakraComponents = {
   [Tag in DOMElements]: ChakraComponent<Tag, {}>
 }
 
-export type HTMLChakraProps<T extends As> = Omit<
+export type HTMLChakraProps<T extends ElementType> = Omit<
   PropsOf<T>,
   "ref" | keyof StyleProps
 > &
-  ChakraProps & { as?: As }
+  ChakraProps &
+  AsChildProps &
+  AsProps

packages/components/src/system/system.types.tsx

@@ -4,6 +4,7 @@ import type {
   SystemStyleObject,
 } from "@chakra-ui/styled-system"
 import type { Interpolation } from "@emotion/react"
+import { ElementType } from "react"
 
 export interface ChakraProps extends SystemProps {
   /**
@@ -26,14 +27,23 @@ export interface ChakraProps extends SystemProps {
   css?: Interpolation<{}>
 }
 
-export type As = React.ElementType
+export interface AsChildProps {
+  asChild?: boolean
+}
+
+export interface AsProps<T extends ElementType = ElementType> {
+  /**
+   * @deprecated - This prop is deprecated and will be removed in a future release
+   * Switch to the `asChild` prop instead
+   */
+  as?: T
+}
 
 /**
  * Extract the props of a React element or component
  */
-export type PropsOf<T extends As> = React.ComponentPropsWithoutRef<T> & {
-  as?: As
-}
+export type PropsOf<T extends ElementType> = React.ComponentPropsWithoutRef<T> &
+  AsProps
 
 export type OmitCommonProps<
   Target,
@@ -56,18 +66,21 @@ export type MergeWithAs<
   ComponentProps extends object,
   AsProps extends object,
   AdditionalProps extends object = {},
-  AsComponent extends As = As,
+  AsComponent extends ElementType = ElementType,
 > = (
   | RightJoinProps<ComponentProps, AdditionalProps>
   | RightJoinProps<AsProps, AdditionalProps>
 ) & {
   as?: AsComponent
 }
 
-export type ComponentWithAs<Component extends As, Props extends object = {}> = {
-  <AsComponent extends As = Component>(
+export type ComponentWithAs<
+  Component extends ElementType,
+  Props extends object = {},
+> = {
+  <AsComponent extends ElementType = Component>(
     props: MergeWithAs<
-      React.ComponentProps<Component>,
+      React.ComponentProps<Component> & AsChildProps,
       React.ComponentProps<AsComponent>,
       Props,
       AsComponent
@@ -81,5 +94,5 @@ export type ComponentWithAs<Component extends As, Props extends object = {}> = {
   id?: string
 }
 
-export interface ChakraComponent<T extends As, P extends object = {}>
+export interface ChakraComponent<T extends ElementType, P extends object = {}>
   extends ComponentWithAs<T, Assign<ChakraProps, P>> {}