From b2f29813770eb72cebb8eeb5dee2714c69d9a6de Mon Sep 17 00:00:00 2001
From: zernonia <zernonia@gmail.com>
Date: Wed, 19 Jun 2024 08:53:16 +0800
Subject: [PATCH] refactor: transform tw prefix by specific case

---
 .../utils/transformers/transform-tw-prefix.ts | 34 ++++++++++++++-----
 .../transform-tw-prefix.test.ts.snap          |  5 +++
 .../test/utils/transform-tw-prefix.test.ts    |  5 +++
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/packages/cli/src/utils/transformers/transform-tw-prefix.ts b/packages/cli/src/utils/transformers/transform-tw-prefix.ts
index 6059a4e5..dfc7a8a1 100644
--- a/packages/cli/src/utils/transformers/transform-tw-prefix.ts
+++ b/packages/cli/src/utils/transformers/transform-tw-prefix.ts
@@ -1,14 +1,13 @@
 import type { CodemodPlugin } from 'vue-metamorph'
 import { splitClassName } from './transform-css-vars'
 import type { TransformOpts } from '.'
-import type { Config } from '@/src/utils/get-config'
 
 export function transformTwPrefix(opts: TransformOpts): CodemodPlugin {
   return {
     type: 'codemod',
     name: 'add prefix to tailwind classes',
 
-    transform({ scriptASTs, sfcAST, utils: { traverseScriptAST, traverseTemplateAST } }) {
+    transform({ scriptASTs, sfcAST, utils: { traverseScriptAST, traverseTemplateAST, astHelpers } }) {
       let transformCount = 0
       const { config } = opts
 
@@ -17,13 +16,32 @@ export function transformTwPrefix(opts: TransformOpts): CodemodPlugin {
 
       for (const scriptAST of scriptASTs) {
         traverseScriptAST(scriptAST, {
-          visitLiteral(path) {
-            if (path.parent.value.type !== 'ImportDeclaration' && typeof path.node.value === 'string') {
-            // mutate the node
-              path.node.value = applyPrefix(path.node.value, config.tailwind.prefix)
-              transformCount++
-            }
+          visitCallExpression(path) {
+            if (path.node.callee.type === 'Identifier' && path.node.callee.name === 'cva') {
+              const nodes = path.node.arguments
+              nodes.forEach((node) => {
+                // cva(base, ...)
+                if (node.type === 'Literal' && typeof node.value === 'string') {
+                  node.value = applyPrefix(node.value, config.tailwind.prefix)
+                  transformCount++
+                }
 
+                else if (node.type === 'ObjectExpression') {
+                  node.properties.forEach((node) => {
+                    // cva(..., { variants: { ... } })
+                    if (node.type === 'Property' && node.key.type === 'Identifier' && node.key.name === 'variants') {
+                      const nodes = astHelpers.findAll(node, { type: 'Literal' })
+                      nodes.forEach((node) => {
+                        if (typeof node.value === 'string') {
+                          node.value = applyPrefix(node.value, config.tailwind.prefix)
+                          transformCount++
+                        }
+                      })
+                    }
+                  })
+                }
+              })
+            }
             return this.traverse(path)
           },
         })
diff --git a/packages/cli/test/utils/__snapshots__/transform-tw-prefix.test.ts.snap b/packages/cli/test/utils/__snapshots__/transform-tw-prefix.test.ts.snap
index 1e4edcd1..795b6230 100644
--- a/packages/cli/test/utils/__snapshots__/transform-tw-prefix.test.ts.snap
+++ b/packages/cli/test/utils/__snapshots__/transform-tw-prefix.test.ts.snap
@@ -2,6 +2,8 @@
 
 exports[`transform tailwind prefix 1`] = `
 "import { cva } from 'class-variance-authority'
+export { default as Button } from './Button.vue'
+export type { ButtonType } from './Button.vue'
 
 export const testVariants = cva(
   'tw-bg-background hover:tw-bg-muted tw-text-primary-foreground sm:focus:tw-text-accent-foreground',
@@ -14,6 +16,9 @@ export const testVariants = cva(
         default: 'tw-h-10 tw-px-4 tw-py-2',
       },
     },
+    defaultVariants: {
+      variant: 'default',
+    },
   },
 )
 "
diff --git a/packages/cli/test/utils/transform-tw-prefix.test.ts b/packages/cli/test/utils/transform-tw-prefix.test.ts
index e1cad783..e01c533d 100644
--- a/packages/cli/test/utils/transform-tw-prefix.test.ts
+++ b/packages/cli/test/utils/transform-tw-prefix.test.ts
@@ -7,6 +7,8 @@ it('transform tailwind prefix', async () => {
     await transform({
       filename: 'test.ts',
       raw: `import { cva } from "class-variance-authority" 
+      export { default as Button } from "./Button.vue"
+      export type { ButtonType } from "./Button.vue"
       
       export const testVariants = cva(
         'bg-background hover:bg-muted text-primary-foreground sm:focus:text-accent-foreground',
@@ -19,6 +21,9 @@ it('transform tailwind prefix', async () => {
               default: 'h-10 px-4 py-2',
             },
           },
+          defaultVariants: {
+            variant: 'default',
+          },
         },
       )`,
       config: {