From 04eed56f70000ffcd62dacc2febe4a2285a75d3d Mon Sep 17 00:00:00 2001
From: Daniel Tijerina <tijerina.daniel@heb.com>
Date: Thu, 29 Apr 2021 22:06:48 -0500
Subject: [PATCH] merge withConfig arguments to allow for shouldForwardProp

---
 .prettierrc                                   |  3 +-
 src/visitors/displayNameAndId.js              | 63 ++++++++++++++++++-
 .../add-identifier-and-display-name/code.js   | 15 +++++
 .../add-identifier-and-display-name/output.js | 19 ++++++
 4 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/.prettierrc b/.prettierrc
index 36301bc5..3d1f5e5c 100644
--- a/.prettierrc
+++ b/.prettierrc
@@ -1,5 +1,6 @@
 {
   "semi": false,
   "singleQuote": true,
-  "trailingComma": "es5"
+  "trailingComma": "es5",
+  "arrowParens": "avoid"
 }
diff --git a/src/visitors/displayNameAndId.js b/src/visitors/displayNameAndId.js
index 42c1c1ff..5f60a9e9 100644
--- a/src/visitors/displayNameAndId.js
+++ b/src/visitors/displayNameAndId.js
@@ -17,6 +17,7 @@ const addConfig = t => (path, displayName, componentId) => {
   }
 
   const withConfigProps = []
+
   if (displayName) {
     withConfigProps.push(
       t.objectProperty(
@@ -34,6 +35,34 @@ const addConfig = t => (path, displayName, componentId) => {
     )
   }
 
+  const existingConfig = getExistingConfig(t)(path)
+
+  if (
+    existingConfig &&
+    existingConfig.arguments.length &&
+    !existingConfig.arguments[0].properties.some(prop =>
+      ['displayName', 'componentId'].includes(prop.key.name)
+    )
+  ) {
+    existingConfig.arguments[0].properties.push(...withConfigProps)
+    return
+  }
+
+  if (
+    path.node.callee &&
+    t.isMemberExpression(path.node.callee.callee) &&
+    path.node.callee.callee.property &&
+    path.node.callee.callee.property.name &&
+    path.node.callee.callee.property.name == 'withConfig' &&
+    path.node.callee.arguments.length &&
+    !path.node.callee.arguments[0].properties.some(prop =>
+      ['displayName', 'componentId'].includes(prop.key.name)
+    )
+  ) {
+    path.node.callee.arguments[0].properties.push(...withConfigProps)
+    return
+  }
+
   if (path.node.tag) {
     // Replace x`...` with x.withConfig({ })`...`
     path.node.tag = t.callExpression(
@@ -53,6 +82,29 @@ const addConfig = t => (path, displayName, componentId) => {
   }
 }
 
+const getExistingConfig = t => path => {
+  if (
+    path.node.callee &&
+    t.isMemberExpression(path.node.callee.callee) &&
+    path.node.callee.callee.property &&
+    path.node.callee.callee.property.name &&
+    path.node.callee.callee.property.name == 'withConfig'
+  ) {
+    return path.node.callee
+  }
+
+  if (
+    path.node.callee &&
+    t.isMemberExpression(path.node.callee.callee) &&
+    path.node.callee.callee.object &&
+    path.node.callee.callee.object.callee &&
+    path.node.callee.callee.object.callee.property &&
+    path.node.callee.callee.object.callee.property.name === 'withConfig'
+  ) {
+    return path.node.callee.callee.object
+  }
+}
+
 const getBlockName = file => {
   const name = path.basename(
     file.opts.filename,
@@ -155,7 +207,16 @@ export default t => (path, state) => {
           t.isMemberExpression(path.node.callee.callee) &&
           path.node.callee.callee.property &&
           path.node.callee.callee.property.name &&
-          path.node.callee.callee.property.name !== 'withConfig')
+          path.node.callee.callee.property.name !== 'withConfig') ||
+        // styled(x).withConfig({})
+        (isStyled(t)(path.node.callee, state) &&
+          t.isMemberExpression(path.node.callee.callee) &&
+          path.node.callee.callee.property &&
+          path.node.callee.callee.property.name &&
+          path.node.callee.callee.property.name == 'withConfig' &&
+          !path.node.callee.arguments[0].properties.some((prop) =>
+            ['displayName', 'componentId'].includes(prop.key.name)
+          ))
   ) {
     const displayName =
       useDisplayName(state) &&
diff --git a/test/fixtures/add-identifier-and-display-name/code.js b/test/fixtures/add-identifier-and-display-name/code.js
index a451c72e..2d514389 100644
--- a/test/fixtures/add-identifier-and-display-name/code.js
+++ b/test/fixtures/add-identifier-and-display-name/code.js
@@ -13,3 +13,18 @@ const WrappedComponent3 = styled(Inner)({})
 const WrappedComponent4 = styled(Inner).attrs(() => ({ something: 'else' }))({})
 const WrappedComponent5 = styled.div.attrs(() => ({ something: 'else' }))({})
 const WrappedComponent6 = styled.div.attrs(() => ({ something: 'else' }))``
+const WrappedComponent7 = styled.div.withConfig({
+  shouldForwardProp: () => {},
+})({})
+
+const WrappedComponent8 = styled.div
+  .withConfig({
+    shouldForwardProp: () => {},
+  })
+  .attrs(() => ({ something: 'else' }))({})
+
+const WrappedComponent9 = styled.div
+  .attrs(() => ({ something: 'else' }))
+  .withConfig({
+    shouldForwardProp: () => {},
+  })({})
diff --git a/test/fixtures/add-identifier-and-display-name/output.js b/test/fixtures/add-identifier-and-display-name/output.js
index 1b3376b3..2d81105c 100644
--- a/test/fixtures/add-identifier-and-display-name/output.js
+++ b/test/fixtures/add-identifier-and-display-name/output.js
@@ -51,3 +51,22 @@ const WrappedComponent6 = styled.div.attrs(() => ({
   displayName: "WrappedComponent6",
   componentId: "sc-1cza72q-10"
 })``;
+const WrappedComponent7 = styled.div.withConfig({
+  shouldForwardProp: () => {},
+  displayName: "WrappedComponent7",
+  componentId: "sc-1cza72q-11"
+})({});
+const WrappedComponent8 = styled.div.withConfig({
+  shouldForwardProp: () => {},
+  displayName: "WrappedComponent8",
+  componentId: "sc-1cza72q-12"
+}).attrs(() => ({
+  something: 'else'
+}))({});
+const WrappedComponent9 = styled.div.attrs(() => ({
+  something: 'else'
+})).withConfig({
+  shouldForwardProp: () => {},
+  displayName: "WrappedComponent9",
+  componentId: "sc-1cza72q-13"
+})({});