fix(sc): fix withComponent usage

Author: gregberge

packages/shared/core/styled-engine/emotion.js

@@ -63,14 +63,14 @@ function patchStyledComponent(StyledComponent) {
   StyledComponent.withComponent = (component, options) => {
     // eslint-disable-next-line no-underscore-dangle
     const BaseComponent = StyledComponent.__emotion_base
-    return Object.assign(
-      patchStyledComponent(
+    return patchStyledComponent(
+      Object.assign(
         withComponent(
           props => <BaseComponent as={component} {...props} />,
           options,
         ),
+        StyledComponent,
       ),
-      StyledComponent,
     )
   }
 

packages/shared/core/styled-engine/util.js

@@ -1,8 +1,31 @@
-function patchStyledComponent(StyledComponent, defaultAs) {
-  const { render } = StyledComponent
-  StyledComponent.withComponent = component =>
-    patchStyledComponent({ ...StyledComponent }, component)
-  StyledComponent.render = ({ as = defaultAs, ...props }, ref) =>
+/* eslint-disable no-underscore-dangle */
+import React from 'react'
+
+function patchStyledComponent(StyledComponent) {
+  if (!StyledComponent.target.__smoothUIComponent) {
+    return StyledComponent
+  }
+  const {
+    render,
+    withComponent,
+    defaultProps,
+    propTypes,
+    componentStyle,
+  } = StyledComponent
+  StyledComponent.withComponent = (component, ...args) => {
+    const Target = StyledComponent.target
+    const NewTarget = props => <Target as={component} {...props} />
+    // eslint-disable-next-line no-underscore-dangle
+    NewTarget.__smoothUIComponent = true
+    return patchStyledComponent(
+      Object.assign(withComponent(NewTarget, ...args), {
+        defaultProps,
+        propTypes,
+        componentStyle,
+      }),
+    )
+  }
+  StyledComponent.render = ({ as, ...props }, ref) =>
     render({ forwardedAs: as, ...props }, ref)
   return StyledComponent
 }
@@ -21,5 +44,7 @@ function wrapCreateStyledComponent(createStyledComponent) {
 }
 
 export function wrapStyled(styled) {
-  return component => wrapCreateStyledComponent(styled(component))
+  const newStyled = component => wrapCreateStyledComponent(styled(component))
+  Object.assign(newStyled, styled)
+  return newStyled
 }

packages/shared/core/utils/createComponent.js

@@ -71,6 +71,9 @@ function createComponent(getConfig) {
   const RefComponent = React.forwardRef(forwardRef)
   RefComponent.displayName = InnerComponent.displayName
 
+  // eslint-disable-next-line no-underscore-dangle
+  RefComponent.__smoothUIComponent = true
+
   const StyledComponent = styled(RefComponent)`
     ${style};
     ${applySystem && applySystem(system)};