From 3bc02234b3fa07ccf8e55025c6bd66e3639c9ecb Mon Sep 17 00:00:00 2001
From: JYC <jycouet@gmail.com>
Date: Sun, 10 Dec 2023 15:59:30 +0100
Subject: [PATCH] Fix/too many imports removed (#540)

* :white_check_mark: UPDATE: failing test

* :bug: FIX: fix

* :construction: UPDATE: ts (rework later)

* add changeset (#539)
---
 .changeset/cool-dolls-yawn.md                 |  5 ++
 .../src/lib/transformDecorator.spec.ts        | 50 +++++++++++++++++
 .../src/lib/transformDecorator.ts             | 54 ++++++++++++-------
 3 files changed, 90 insertions(+), 19 deletions(-)
 create mode 100644 .changeset/cool-dolls-yawn.md

diff --git a/.changeset/cool-dolls-yawn.md b/.changeset/cool-dolls-yawn.md
new file mode 100644
index 00000000..fe90ed26
--- /dev/null
+++ b/.changeset/cool-dolls-yawn.md
@@ -0,0 +1,5 @@
+---
+'vite-plugin-striper': patch
+---
+
+fix import in array is now detected
diff --git a/packages/vite-plugin-striper/src/lib/transformDecorator.spec.ts b/packages/vite-plugin-striper/src/lib/transformDecorator.spec.ts
index 166e85a0..a7e8ee5d 100644
--- a/packages/vite-plugin-striper/src/lib/transformDecorator.spec.ts
+++ b/packages/vite-plugin-striper/src/lib/transformDecorator.spec.ts
@@ -194,4 +194,54 @@ export class TasksController {
     }
   `)
   })
+
+  it('should strip just the right things', async () => {
+    const code = `import { Allow, BackendMethod, Entity, Fields, Validators } from 'remult'
+
+    @Entity<User2>('userstest', {
+      allowApiCrud: Allow.authenticated,
+    })
+    export class User2 {
+      @Fields.uuid()
+      id = ''
+    
+      @Fields.string({
+        validate: [Validators.required, Validators.uniqueOnBackend],
+      })
+      email = ''
+    
+      @BackendMethod({ allowed: Allow.everyone })
+      async testMethod() {
+        console.log('hello')
+      }
+    }
+	`
+
+    const transformed = await transformDecorator(code, ['BackendMethod'])
+
+    expect(transformed).toMatchInlineSnapshot(`
+      {
+        "code": "import { Allow, BackendMethod, Entity, Fields } from \\"remult\\";
+
+      @Entity<User2>(\\"userstest\\", {
+          allowApiCrud: Allow.authenticated
+      })
+      export class User2 {
+          @Fields.uuid()
+          id = \\"\\";
+
+          @Fields.string({
+              validate: [Validators.required, Validators.uniqueOnBackend]
+          })
+          email = \\"\\";
+
+          @BackendMethod({
+              allowed: Allow.everyone
+          })
+          async testMethod() {}
+      }",
+        "transformed": true,
+      }
+    `)
+  })
 })
diff --git a/packages/vite-plugin-striper/src/lib/transformDecorator.ts b/packages/vite-plugin-striper/src/lib/transformDecorator.ts
index 9db05c19..0a7f7c45 100644
--- a/packages/vite-plugin-striper/src/lib/transformDecorator.ts
+++ b/packages/vite-plugin-striper/src/lib/transformDecorator.ts
@@ -86,30 +86,22 @@ export const transformDecorator = async (code: string, decorators_to_strip: stri
 
           this.traverse(path)
         },
+
         visitClassDeclaration(path) {
+          // Capture identifiers in class decorators
           // @ts-ignore
-          if (path.node.decorators) {
-            // @ts-ignore
-            const vals = path.node.decorators.map(
-              // @ts-ignore
-              a => a.expression.callee.name,
-            )
-            // @ts-ignore
-            usedIdentifiersInCode.add(...vals)
-          }
+          ;(path.node.decorators || []).forEach(decorator => {
+            extractIdentifiersFromExpression(decorator.expression, usedIdentifiersInCode)
+          })
 
+          // Capture identifiers in class methods and properties
           path.node.body.body.forEach(element => {
-            if (element.type === 'ClassProperty') {
+            if (element.type === 'ClassMethod' || element.type === 'ClassProperty') {
+              // Capture identifiers in element decorators
               // @ts-ignore
-              const vals = element.decorators
-                // @ts-ignore
-                .map(d => d.expression)
-                // @ts-ignore
-                .map(e => e.callee)
-                // @ts-ignore
-                .map(f => f.object.name)
-              // @ts-ignore
-              usedIdentifiersInCode.add(...vals)
+              ;(element.decorators || []).forEach(decorator => {
+                extractIdentifiersFromExpression(decorator.expression, usedIdentifiersInCode)
+              })
             }
           })
 
@@ -143,3 +135,27 @@ export const transformDecorator = async (code: string, decorators_to_strip: stri
     return { code, transformed: false }
   }
 }
+
+// Helper function to extract identifiers from an expression
+// @ts-ignore
+function extractIdentifiersFromExpression(expression, identifierSet) {
+  if (!expression) return
+
+  if (expression.type === 'Identifier') {
+    identifierSet.add(expression.name)
+  } else if (expression.type === 'MemberExpression') {
+    extractIdentifiersFromExpression(expression.object, identifierSet)
+    extractIdentifiersFromExpression(expression.property, identifierSet)
+  } else if (expression.type === 'CallExpression') {
+    extractIdentifiersFromExpression(expression.callee, identifierSet)
+    // @ts-ignore
+    expression.arguments.forEach(arg => extractIdentifiersFromExpression(arg, identifierSet))
+  } else if (expression.type === 'ArrayExpression') {
+    // Process each element in the array
+    // @ts-ignore
+    expression.elements.forEach(element => {
+      extractIdentifiersFromExpression(element, identifierSet)
+    })
+  }
+  // Add other expression types as needed
+}
-- 
2.51.2