From b0ce8e9166dcc5aba6cd3fbb1558a8f3f25c8a18 Mon Sep 17 00:00:00 2001
From: Akshit Garg <153659573+AkshitGarg24@users.noreply.github.com>
Date: Mon, 23 Sep 2024 01:43:25 +0530
Subject: [PATCH] updated styledictionary/4/asynchronous-api codemod

---
 .../4/asynchronous-api/src/index.ts           | 51 ++++++++++++-------
 1 file changed, 34 insertions(+), 17 deletions(-)

diff --git a/packages/codemods/styledictionary/4/asynchronous-api/src/index.ts b/packages/codemods/styledictionary/4/asynchronous-api/src/index.ts
index 2761aa271..b50af8ae6 100644
--- a/packages/codemods/styledictionary/4/asynchronous-api/src/index.ts
+++ b/packages/codemods/styledictionary/4/asynchronous-api/src/index.ts
@@ -1,9 +1,4 @@
-import type {
-    API,
-    ArrowFunctionExpression,
-    FileInfo,
-    Options,
-} from 'jscodeshift';
+import type { API, FileInfo, Options } from 'jscodeshift';
 
 export default function transform(
     file: FileInfo,
@@ -25,6 +20,32 @@ export default function transform(
 
     const root = j(file.source);
 
+    // Keep track of variables declared using StyleDictionary
+    const styleDictionaryVars: Set<string> = new Set();
+
+    // Identify variables assigned using StyleDictionary.extend or new StyleDictionary
+    root.find(j.VariableDeclarator).forEach((path) => {
+        const init = path.node.init;
+
+        // Handle StyleDictionary.extend() pattern
+        if (
+            j.CallExpression.check(init) &&
+            init.callee.type === 'MemberExpression' &&
+            init.callee.object.name === 'StyleDictionary' &&
+            init.callee.property.name === 'extend'
+        ) {
+            styleDictionaryVars.add(path.node.id.name);
+        }
+
+        // Handle new StyleDictionary() pattern
+        if (
+            j.NewExpression.check(init) &&
+            init.callee.name === 'StyleDictionary'
+        ) {
+            styleDictionaryVars.add(path.node.id.name);
+        }
+    });
+
     // Find all function declarations, function expressions, and arrow functions
     root.find(j.FunctionDeclaration).forEach((path) =>
         addAsyncIfNeeded(path, j),
@@ -38,21 +59,17 @@ export default function transform(
         addAsyncIfNeeded(path, j),
     );
 
-    // Add 'await' before the relevant method calls, except for StyleDictionary.extend()
+    // Add 'await' before the relevant method calls, only for variables declared using StyleDictionary
     asyncMethods.forEach((method) => {
         root.find(j.CallExpression, {
             callee: { property: { name: method } },
         }).forEach((path) => {
-            // Check if the method is already awaited
-            if (!j.AwaitExpression.check(path.parent.node)) {
-                // Skip if it's StyleDictionary.extend()
-                if (
-                    !(
-                        path.node.callee.object &&
-                        path.node.callee.object.name === 'StyleDictionary' &&
-                        path.node.callee.property.name === 'extend'
-                    )
-                ) {
+            const objectName = path.node.callee.object.name;
+
+            // Ensure the method is called on a variable initialized with StyleDictionary
+            if (styleDictionaryVars.has(objectName)) {
+                // Check if the method is already awaited
+                if (!j.AwaitExpression.check(path.parent.node)) {
                     j(path).replaceWith(j.awaitExpression(path.node));
                 }
             }