feat(to-promise-all): new command, close #5

Author: antfu

README.md

@@ -326,6 +326,34 @@ const quxx = `${qux}quxx`
 const foo = `foo`; const baz = 'baz'; const qux = `qux`
 ```
 
+### `to-promise-all`
+
+Convert multiple `await` statements to `await Promise.all()`.
+
+Triggers:
+- `/// to-promise-all`
+- `/// 2pa
+
+```js
+/// to-promise-all
+const foo = await getFoo()
+const { bar, baz } = await getBar()
+```
+
+Will be converted to:
+
+```js
+const [
+  foo,
+  { bar, baz },
+] = await Promise.all([
+  getFoo(),
+  getBar(),
+])
+```
+
+This command will try to search all continuous declarations with `await` and convert them to a single `await Promise.all()` call.
+
 ## Custom Commands
 
 It's also possible to define your custom commands.

src/commands/index.ts

@@ -7,6 +7,7 @@ import { toDynamicImport } from './to-dynamic-import'
 import { toStringLiteral } from './to-string-literal'
 import { toTemplateLiteral } from './to-template-literal'
 import { inlineArrow } from './inline-arrow'
+import { toPromiseAll } from './to-promise-all'
 
 // @keep-sorted
 export {
@@ -17,6 +18,7 @@ export {
   toForEach,
   toForOf,
   toFunction,
+  toPromiseAll,
   toStringLiteral,
   toTemplateLiteral,
 }
@@ -30,6 +32,7 @@ export const builtinCommands = [
   toForEach,
   toForOf,
   toFunction,
+  toPromiseAll,
   toStringLiteral,
   toTemplateLiteral,
 ]

src/commands/to-promise-all.test.ts

@@ -0,0 +1,149 @@
+import { toPromiseAll as command } from './to-promise-all'
+import { d, run } from './_test-utils'
+
+run(
+  command,
+  // Program level
+  {
+    code: d`
+      /// to-promise-all
+      const a = await foo()
+      const b = await bar()
+    `,
+    output: d`
+      const [
+      a,
+      b,
+      ] = await Promise.all([
+      foo(),
+      bar(),
+      ])
+    `,
+    errors: ['command-removal', 'command-fix'],
+  },
+  // Function declaration
+  {
+    filename: 'index.ts',
+    code: d`
+      async function fn() {
+        /// to-promise-all
+        const a = await foo()
+        const b = await bar()
+      }
+    `,
+    output: d`
+      async function fn() {
+        const [
+        a,
+        b,
+        ] = await Promise.all([
+        foo(),
+        bar(),
+        ] as const)
+      }
+    `,
+    errors: ['command-removal', 'command-fix'],
+  },
+  // If Statement
+  {
+    code: d`
+      if (true) {
+        /// to-promise-all
+        const a = await foo()
+          .then(() => {})
+        const b = await import('bar').then(m => m.default)
+      }
+    `,
+    output: d`
+      if (true) {
+        const [
+        a,
+        b,
+        ] = await Promise.all([
+        foo()
+          .then(() => {}),
+        import('bar').then(m => m.default),
+        ])
+      }
+    `,
+    errors: ['command-removal', 'command-fix'],
+  },
+  // Mixed declarations
+  {
+    code: d`
+      on('event', async () => {
+        /// to-promise-all
+        let a = await foo()
+          .then(() => {})
+        const { foo, bar } = await import('bar').then(m => m.default)
+        const b = await baz(), c = await qux(), d = foo()
+      })
+    `,
+    output: d`
+      on('event', async () => {
+        let [
+        a,
+        { foo, bar },
+        b,
+        c,
+        d,
+        ] = await Promise.all([
+        foo()
+          .then(() => {}),
+        import('bar').then(m => m.default),
+        baz(),
+        qux(),
+        foo(),
+        ])
+      })
+    `,
+    errors: ['command-removal', 'command-fix'],
+  },
+  // Await expressions
+  {
+    code: d`
+      /// to-promise-all
+      const a = await bar()
+      await foo()
+      const b = await baz()
+      doSomething()
+      const nonTarget = await qux()
+    `,
+    output: d`
+      const [
+      a,
+      /* discarded */,
+      b,
+      ] = await Promise.all([
+      bar(),
+      foo(),
+      baz(),
+      ])
+      doSomething()
+      const nonTarget = await qux()
+    `,
+    errors: ['command-removal', 'command-fix'],
+  },
+  // Should stop on first non-await expression
+  {
+    code: d`
+      /// to-promise-all
+      const a = await bar()
+      let b = await foo()
+      let c = baz()
+      const d = await qux()
+    `,
+    output: d`
+      let [
+      a,
+      b,
+      ] = await Promise.all([
+      bar(),
+      foo(),
+      ])
+      let c = baz()
+      const d = await qux()
+    `,
+    errors: ['command-removal', 'command-fix'],
+  },
+)

src/commands/to-promise-all.ts

@@ -0,0 +1,97 @@
+import type { Command, Tree } from '../types'
+
+type TargetNode = Tree.VariableDeclaration | Tree.ExpressionStatement
+type TargetDeclarator = Tree.VariableDeclarator | Tree.AwaitExpression
+
+export const toPromiseAll: Command = {
+  name: 'to-promise-all',
+  match: /^[\/@:]\s*(?:to-|2)(?:promise-all|pa)$/,
+  action(ctx) {
+    const parent = ctx.getParentBlock()
+    const nodeStart = ctx.findNodeBelow(isTarget) as TargetNode
+    let nodeEnd: Tree.Node = nodeStart
+    if (!nodeStart)
+      return ctx.reportError('Unable to find variable declaration')
+    if (!parent.body.includes(nodeStart))
+      return ctx.reportError('Variable declaration is not in the same block')
+
+    function isTarget(node: Tree.Node): node is TargetNode {
+      if (node.type === 'VariableDeclaration')
+        return node.declarations.some(declarator => declarator.init?.type === 'AwaitExpression')
+      else if (node.type === 'ExpressionStatement')
+        return node.expression.type === 'AwaitExpression'
+      return false
+    }
+
+    function getDeclarators(node: TargetNode): TargetDeclarator[] {
+      if (node.type === 'VariableDeclaration')
+        return node.declarations
+      if (node.expression.type === 'AwaitExpression')
+        return [node.expression]
+      return []
+    }
+
+    let declarationType = 'const'
+    const declarators: TargetDeclarator[] = []
+    for (let i = parent.body.indexOf(nodeStart); i < parent.body.length; i++) {
+      const node = parent.body[i]
+      if (isTarget(node)) {
+        declarators.push(...getDeclarators(node))
+        nodeEnd = node
+        if (node.type === 'VariableDeclaration' && node.kind !== 'const')
+          declarationType = 'let'
+      }
+      else {
+        break
+      }
+    }
+
+    ctx.removeComment()
+    ctx.report({
+      loc: {
+        start: nodeStart.loc.start,
+        end: nodeEnd.loc.end,
+      },
+      message: 'Convert to `await Promise.all`',
+      fix(fixer) {
+        const lineIndent = ctx.getIndentOfLine(nodeStart.loc.start.line)
+        const isTs = ctx.context.filename.match(/\.[mc]?tsx?$/)
+
+        function unwrapAwait(node: Tree.Node | null) {
+          if (node?.type === 'AwaitExpression')
+            return node.argument
+          return node
+        }
+
+        function getId(declarator: TargetDeclarator) {
+          if (declarator.type === 'AwaitExpression')
+            return '/* discarded */'
+          return ctx.getTextOf(declarator.id)
+        }
+
+        function getInit(declarator: TargetDeclarator) {
+          if (declarator.type === 'AwaitExpression')
+            return ctx.getTextOf(declarator.argument)
+          return ctx.getTextOf(unwrapAwait(declarator.init))
+        }
+
+        const str = [
+          `${declarationType} [`,
+          ...declarators
+            .map(declarator => `${getId(declarator)},`),
+          '] = await Promise.all([',
+          ...declarators
+            .map(declarator => `${getInit(declarator)},`),
+          isTs ? '] as const)' : '])',
+        ]
+          .map((line, idx) => idx ? lineIndent + line : line)
+          .join('\n')
+
+        return fixer.replaceTextRange([
+          nodeStart.range[0],
+          nodeEnd.range[1],
+        ], str)
+      },
+    })
+  },
+}

src/context.ts

@@ -204,4 +204,27 @@ export class CommandContext {
       ? result
       : result[0]
   }
+
+  /**
+   * Get the parent block of the triggering comment
+   */
+  getParentBlock(): Tree.BlockStatement | Tree.Program {
+    const node = this.source.getNodeByRangeIndex(this.comment.range[0])
+    if (node?.type === 'BlockStatement') {
+      if (this.source.getCommentsInside(node).includes(this.comment))
+        return node
+    }
+    if (node)
+      console.warn(`Expected BlockStatement, got ${node.type}. This is probably an internal bug.`)
+    return this.source.ast
+  }
+
+  /**
+   * Get indent string of a specific line
+   */
+  getIndentOfLine(line: number): string {
+    const lineStr = this.source.getLines()[line - 1] || ''
+    const match = lineStr.match(/^\s*/)
+    return match ? match[0] : ''
+  }
 }