fix: fix withComponent

Author: gregberge

src/Button.js

@@ -3,7 +3,7 @@ import PropTypes from 'prop-types'
 import classNames from 'classnames'
 import styled, { css } from 'styled-components'
 import handleRef from './internal/handleRef'
-import getWithComponent from './internal/getWithComponent'
+import setWithComponent from './internal/setWithComponent'
 import * as defaultTheme from './style/defaultTheme'
 import { th, mixin } from './utils'
 
@@ -88,7 +88,7 @@ Button.defaultProps = {
   theme: defaultTheme,
 }
 
-Button.withComponent = getWithComponent(Button, ButtonComponent)
+setWithComponent(Button, ButtonComponent)
 
 /** @component */
 export default Button

src/Button.md

@@ -56,3 +56,20 @@ Disable using `disabled` prop.
 ```js
 <Button disabled>Disabled</Button>
 ```
+
+### Extend and use with other components
+
+```js
+const BorderedButton = Button.extend`
+  border: 2px solid black;
+`
+const BorderedLinkButton = BorderedButton.withComponent('a')
+;<div>
+  <span style={{ margin: '5px' }}>
+    <BorderedLinkButton variant="primary">Primary</BorderedLinkButton>
+  </span>
+  <span style={{ margin: '5px' }}>
+    <BorderedLinkButton variant="secondary">Secondary</BorderedLinkButton>
+  </span>
+</div>
+```

src/Input.js

@@ -3,7 +3,7 @@ import classNames from 'classnames'
 import PropTypes from 'prop-types'
 import styled from 'styled-components'
 import handleRef from './internal/handleRef'
-import getWithComponent from './internal/getWithComponent'
+import setWithComponent from './internal/setWithComponent'
 import * as defaultTheme from './style/defaultTheme'
 import { th, mixin } from './utils'
 
@@ -113,7 +113,7 @@ Input.defaultProps = {
   theme: defaultTheme,
 }
 
-Input.withComponent = getWithComponent(Input, InputComponent)
+setWithComponent(Input, InputComponent)
 
 /** @component */
 export default Input

src/internal/getWithComponent.js

@@ -0,0 +1,29 @@
+/* eslint-disable no-param-reassign */
+import React from 'react'
+
+const setWithComponent = (StyledComponent, BaseComponent) => {
+  const originalExtend = StyledComponent.extend.bind(this)
+  Object.defineProperty(StyledComponent, 'extend', {
+    get() {
+      return (...args) => {
+        const NewStyledComponent = originalExtend(...args)
+        setWithComponent(NewStyledComponent, BaseComponent)
+        return NewStyledComponent
+      }
+    },
+  })
+
+  const originalWithComponent = StyledComponent.withComponent.bind(
+    StyledComponent,
+  )
+  StyledComponent.withComponent = component => {
+    const NewStyledComponent = originalWithComponent(props => (
+      <BaseComponent component={component} {...props} />
+    ))
+
+    setWithComponent(NewStyledComponent, BaseComponent)
+    return NewStyledComponent
+  }
+}
+
+export default setWithComponent