feat: add astx.flatMap()

codemodsquad · Mar 30, 2024 · 3dc5199 · 3dc5199
1 parent 0e61bde
commit 3dc5199
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/src/Astx.ts b/src/Astx.ts
@@ -257,6 +257,17 @@ export default class Astx extends ExtendableProxy implements Iterable<Astx> {
     return result
   }
 
+  flatMap<T>(
+    iteratee: (astx: Astx, index: number, parent: Astx) => T | T[]
+  ): T[] {
+    const result: any[] = []
+    let index = 0
+    for (const astx of this) {
+      result.push(iteratee(astx, index++, this))
+    }
+    return result.flat()
+  }
+
   at(index: number): Astx {
     return new Astx(
       this.context,

diff --git a/test/astx/splitImports.ts b/test/astx/splitImports.ts
@@ -6,11 +6,12 @@ import { mapValues, map } from 'lodash'
 import { fromPairs } from 'lodash'
 `
 
-export function astx({ astx, statement }: TransformOptions): void {
-  astx.find`import { $$imports } from 'lodash'`.replace(({ $$imports }) =>
-    $$imports.map(
-      (imp) => statement`import ${imp.code} from 'lodash/${imp.code}'`
-    )
+export function astx({ astx }: TransformOptions): void {
+  astx.find`import { $$imports } from 'lodash'`.replace(
+    ({ $$imports }, parse) =>
+      $$imports.flatMap(
+        (imp) => parse`import ${imp.code} from 'lodash/${imp.code}'`
+      )
   )
 }