From 8d72fae8c4ba68c8d94c52b5b2b325cd003e045a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krist=C3=B3f=20Poduszl=C3=B3?= <kripod@protonmail.com>
Date: Sun, 16 Oct 2022 14:27:15 +0200
Subject: [PATCH] feat: support properties and variable declarators

---
 src/index.js  | 37 ++++++++++++++++++++++++++++++++-----
 tests/test.js | 23 +++++++++++++++++++----
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/src/index.js b/src/index.js
index 5fa72c3..3dd6b39 100644
--- a/src/index.js
+++ b/src/index.js
@@ -398,6 +398,14 @@ function sortTemplateLiteral(node, { env }) {
   return didChange
 }
 
+function trySortIfLiteral(node, { env }) {
+  if (isStringLiteral(node)) {
+    sortStringLiteral(node, { env })
+  } else if (node.type === 'TemplateLiteral') {
+    sortTemplateLiteral(node, { env })
+  }
+}
+
 function transformJavaScript(ast, { env }) {
   visit(ast, {
     JSXAttribute(node) {
@@ -409,15 +417,34 @@ function transformJavaScript(ast, { env }) {
           sortStringLiteral(node.value, { env })
         } else if (node.value.type === 'JSXExpressionContainer') {
           visit(node.value, (node, parent, key) => {
-            if (isStringLiteral(node)) {
-              sortStringLiteral(node, { env })
-            } else if (node.type === 'TemplateLiteral') {
-              sortTemplateLiteral(node, { env })
-            }
+            trySortIfLiteral(node, { env })
           })
         }
       }
     },
+    VariableDeclarator(node) {
+      if (!node.init) {
+        return
+      }
+      if (node.id.name === 'className') {
+        visit(node.init, (node) => {
+          trySortIfLiteral(node, { env })
+        })
+      }
+    },
+    Property(node) {
+      if (
+        (node.key.type === 'Identifier' && node.key.name === 'className') ||
+        (isStringLiteral(node.key) && node.key.value === 'className')
+      ) {
+        visit(node.value, (node) => {
+          trySortIfLiteral(node, { env })
+        })
+      }
+    },
+    ObjectProperty(node) {
+      return this.Property(node)
+    },
   })
 }
 
diff --git a/tests/test.js b/tests/test.js
index 207051a..374cd9f 100644
--- a/tests/test.js
+++ b/tests/test.js
@@ -72,7 +72,7 @@ let css = [
   ['@apply sm:p-0\n   p-0;', '@apply p-0\n   sm:p-0;'],
 ]
 
-let javascript = [
+let jsx = [
   t`;<div class="${yes}" />`,
   t`/* <div class="${no}" /> */`,
   t`// <div class="${no}" />`,
@@ -100,10 +100,25 @@ let javascript = [
     `;<div class={\`p-0 sm:p-0 \${someVar}sm:block flex md:inline\`} />`,
   ],
 ]
-javascript = javascript.concat(
-  javascript.map((test) => test.map((t) => t.replace(/class/g, 'className')))
+jsx = jsx.concat(
+  jsx.map((test) => test.map((t) => t.replace(/class/g, 'className')))
 )
 
+let javascript = [
+  ...jsx,
+  t`var className = '${yes}'`,
+  t`var notClassName = '${no}'`,
+  t`let className = '${yes}'`,
+  t`let notClassName = '${no}'`,
+  t`const className = '${yes}'`,
+  t`const notClassName = '${no}'`,
+  t`let className = \`${yes} \${'${yes}'} \${'${yes}' ? '${yes}' : '${yes}'}\``,
+  t`let obj = { className: \`${yes} \${'${yes}'} \${'${yes}' ? '${yes}' : '${yes}'}\` }`,
+  t`let obj = { className: '${yes}' ? '${yes}' + '${yes}' : '${yes}' }`,
+  t`let obj = { ['className']: \`${yes} \${'${yes}'} \${'${yes}' ? '${yes}' : '${yes}'}\` }`,
+  t`let obj = { ['className']: '${yes}' ? '${yes}' + '${yes}' : '${yes}' }`,
+]
+
 let vue = [
   ...html,
   t`<div :class="'${yes}'"></div>`,
@@ -189,7 +204,7 @@ let tests = {
   'babel-flow': javascript,
   espree: javascript,
   meriyah: javascript,
-  mdx: javascript
+  mdx: jsx
     .filter((test) => !test.find((t) => /^\/\*/.test(t)))
     .map((test) => test.map((t) => t.replace(/^;/, ''))),
   svelte: [