feat: beforeWriteFile supports return Promise

qmhc · Jul 13, 2023 · 5e5b93c · 5e5b93c
1 parent 9c1cfb1
commit 5e5b93c
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 59 deletions.
diff --git a/README.md b/README.md
@@ -280,20 +280,21 @@ export interface PluginOptions {
    *
    * This allows you to transform the path or content
    *
-   * The file will be skipped when the return value `false`
+   * The file will be skipped when the return value `false` or `Promise<false>`
    *
    * @default () => {}
    */
   beforeWriteFile?: (
     filePath: string,
     content: string
-  ) =>
-  | void
-  | false
-  | {
-    filePath?: string,
-    content?: string
-  },
+  ) => MaybePromise<
+    | void
+    | false
+    | {
+      filePath?: string,
+      content?: string
+    }
+  >,
 
   /**
    * Hook called after all declaration files are written

diff --git a/README.zh-CN.md b/README.zh-CN.md
@@ -280,20 +280,21 @@ export interface PluginOptions {
    *
    * 可以在钩子里转换文件路径和文件内容
    *
-   * 当返回 `false` 时会跳过该文件
+   * 当返回 `false` 或 `Promise<false>` 时会跳过该文件
    *
    * @default () => {}
    */
   beforeWriteFile?: (
     filePath: string,
     content: string
-  ) =>
-  | void
-  | false
-  | {
-    filePath?: string,
-    content?: string
-  },
+  ) => MaybePromise<
+    | void
+    | false
+    | {
+      filePath?: string,
+      content?: string
+    }
+  >,
 
   /**
    * 在所有类型文件被写入后调用的钩子

diff --git a/src/plugin.ts b/src/plugin.ts
@@ -22,13 +22,13 @@ import {
   ensureAbsolute,
   ensureArray,
   isNativeObj,
-  isPromise,
   isRegExp,
   normalizePath,
   queryPublicPath,
   removeDirIfEmpty,
   resolve,
-  runParallel
+  runParallel,
+  wrapPromise
 } from './utils'
 
 import type { Alias, Logger } from 'vite'
@@ -320,9 +320,7 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
       }
 
       if (typeof afterDiagnostic === 'function') {
-        const result = afterDiagnostic(diagnostics)
-
-        isPromise(result) && (await result)
+        await wrapPromise(afterDiagnostic(diagnostics))
       }
 
       rootNames.forEach(file => {
@@ -427,7 +425,18 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
       const outDir = outDirs[0]
       const emittedFiles = new Map<string, string>()
 
-      const writeOutput = async (path: string, content: string, outDir: string) => {
+      const writeOutput = async (path: string, content: string, outDir: string, record = true) => {
+        if (typeof beforeWriteFile === 'function') {
+          const result = await wrapPromise(beforeWriteFile(path, content))
+
+          if (result === false) return
+
+          if (result) {
+            path = result.filePath || path
+            content = result.content ?? content
+          }
+        }
+
         path = normalizePath(path)
         const dir = normalizePath(dirname(path))
 
@@ -441,7 +450,7 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
         }
 
         await writeFile(path, content, 'utf-8')
-        emittedFiles.set(path, content)
+        record && emittedFiles.set(path, content)
       }
 
       const service = program.__vue.languageService
@@ -486,18 +495,6 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
           )
           content = cleanVueFileName ? content.replace(/['"](.+)\.vue['"]/g, '"$1"') : content
 
-          if (typeof beforeWriteFile === 'function') {
-            const result = beforeWriteFile(path, content)
-
-            // #110 skip if return false
-            if (result === false) return
-
-            if (result && isNativeObj(result)) {
-              path = result.filePath || path
-              content = result.content ?? content
-            }
-          }
-
           await writeOutput(path, content, outDir)
         }
       )
@@ -519,7 +516,7 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
         const typesPath = types ? resolve(root, types) : resolve(outDir, indexName)
 
         for (const name of entryNames) {
-          let path = multiple ? resolve(outDir, `${name.replace(tsRE, '')}.d.ts`) : typesPath
+          const path = multiple ? resolve(outDir, `${name.replace(tsRE, '')}.d.ts`) : typesPath
 
           if (existsSync(path)) continue
 
@@ -543,17 +540,6 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
             }
           }
 
-          if (typeof beforeWriteFile === 'function') {
-            const result = beforeWriteFile(path, content)
-
-            if (result === false) return
-
-            if (result && isNativeObj(result)) {
-              path = result.filePath ?? path
-              content = result.content ?? content
-            }
-          }
-
           await writeOutput(path, content, outDir)
         }
 
@@ -630,16 +616,14 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
 
           await Promise.all(
             extraOutDirs.map(async outDir => {
-              await writeOutput(resolve(outDir, relativePath), content, outDir)
+              await writeOutput(resolve(outDir, relativePath), content, outDir, false)
             })
           )
         })
       }
 
       if (typeof afterBuild === 'function') {
-        const result = afterBuild()
-
-        isPromise(result) && (await result)
+        await wrapPromise(afterBuild())
       }
 
       bundleDebug('finish')

diff --git a/src/types.ts b/src/types.ts
@@ -199,20 +199,21 @@ export interface PluginOptions {
    *
    * This allows you to transform the path or content
    *
-   * The file will be skipped when the return value `false`
+   * The file will be skipped when the return value `false` or `Promise<false>`
    *
    * @default () => {}
    */
   beforeWriteFile?: (
     filePath: string,
     content: string
-  ) =>
-  | void
-  | false
-  | {
-    filePath?: string,
-    content?: string
-  },
+  ) => MaybePromise<
+    | void
+    | false
+    | {
+      filePath?: string,
+      content?: string
+    }
+  >,
 
   /**
    * Hook called after all declaration files are written

diff --git a/src/utils.ts b/src/utils.ts
@@ -37,6 +37,10 @@ export function isPromise(value: unknown): value is Promise<any> {
   )
 }
 
+export async function wrapPromise<T>(maybePromise: T | Promise<T>) {
+  return isPromise(maybePromise) ? await maybePromise : maybePromise
+}
+
 export function mergeObjects<T extends Record<string, unknown>, U extends Record<string, unknown>>(
   sourceObj: T,
   targetObj: U