feat: collect declared modules when rollup types

fix #240

examples/react/src/index.ts

@@ -2,3 +2,4 @@ import HelloWorld from './components/HelloWorld'
 
 export { HelloWorld }
 export { default as App } from './App'
+export * from './modules'

examples/react/src/modules.ts

@@ -0,0 +1,5 @@
+declare module 'react' {
+  export type Test = { newField: string }
+}
+
+export {}

examples/vue/src/index.ts

@@ -28,6 +28,7 @@ export const user: import('@/src/types').User = {
 
 export type { User } from './types'
 export type { DtsType } from './dts-types'
+export type * from './modules'
 
 export type { AliasType } from '$alias/type'
 

examples/vue/src/main.ts

@@ -3,3 +3,4 @@ export { default as Setup } from '@/components/Setup.vue'
 
 export type { User } from './types'
 export type { DtsType } from './dts-types'
+export type * from './modules'

examples/vue/src/modules.ts

@@ -0,0 +1,5 @@
+declare module 'vue' {
+  export type Test = { newField: string }
+}
+
+export {}

examples/vue/vite.config.ts

@@ -44,7 +44,7 @@ export default defineConfig({
       // include: ['src/index.ts'],
       exclude: ['src/ignore'],
       // staticImport: true,
-      // rollupTypes: true,
+      rollupTypes: true,
       // insertTypesEntry: true,
       compilerOptions: {
         declarationMap: true

src/plugin.ts

@@ -478,6 +478,7 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
 
       const outDir = outDirs[0]
       const emittedFiles = new Map<string, string>()
+      const declareModules: string[] = []
 
       const writeOutput = async (path: string, content: string, outDir: string, record = true) => {
         if (typeof beforeWriteFile === 'function') {
@@ -542,14 +543,17 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
           const baseDir = dirname(filePath)
 
           if (!isMapFile && content) {
-            content = transformCode({
+            const result = transformCode({
               filePath,
               content,
               aliases,
               aliasesExclude,
               staticImport,
               clearPureImport
             })
+
+            content = result.content
+            declareModules.push(...result.declareModules)
           }
 
           filePath = resolve(
@@ -695,9 +699,15 @@ export function dtsPlugin(options: PluginOptions = {}): import('vite').Plugin {
           removeDirIfEmpty(outDir)
           emittedFiles.clear()
 
-          for (const file of rollupFiles) {
-            emittedFiles.set(file, await readFile(file, 'utf-8'))
-          }
+          const declared = declareModules.join('\n')
+
+          await runParallel(cpus().length, [...rollupFiles], async filePath => {
+            await writeOutput(
+              filePath,
+              (await readFile(filePath, 'utf-8')) + (declared ? `\n${declared}\n` : ''),
+              dirname(filePath)
+            )
+          })
 
           bundleDebug('rollup output')
         }

src/transform.ts

@@ -95,6 +95,7 @@ export function transformCode(options: {
 
   const importMap = new Map<string, Set<string>>()
   const usedDefault = new Map<string, string>()
+  const declareModules: string[] = []
 
   let indexCount = 0
 
@@ -214,13 +215,20 @@ export function transformCode(options: {
 
       return false
     }
+
+    if (ts.isModuleDeclaration(node)) {
+      declareModules.push(s.slice(node.pos, node.end + 1))
+    }
   })
 
   importMap.forEach((importSet, libName) => {
     s.prepend(`import { ${Array.from(importSet).join(', ')} } from '${libName}';\n`)
   })
 
-  return s.toString()
+  return {
+    content: s.toString(),
+    declareModules
+  }
 }
 
 export function hasExportDefault(content: string) {

tests/transform.spec.ts

@@ -33,23 +33,23 @@ describe('transform tests', () => {
     })
 
     expect(
-      transformCode(options('let data: import("vexip-ui/lib/tree").InitDataOptions[];'))
+      transformCode(options('let data: import("vexip-ui/lib/tree").InitDataOptions[];')).content
     ).toEqual("import { InitDataOptions } from 'vexip-ui/lib/tree';\nlet data: InitDataOptions[];")
 
     expect(
       transformCode(
         options(
           'declare const _default: import("vue").DefineComponent<{}, {}, {}, {}, {}, import("vue").ComponentOptionsMixin>;\nexport default _default;\n'
         )
-      )
+      ).content
     ).toEqual(
       "import { DefineComponent, ComponentOptionsMixin } from 'vue';\ndeclare const _default: DefineComponent<{}, {}, {}, {}, {}, ComponentOptionsMixin>;\nexport default _default;\n"
     )
 
     expect(
       transformCode(
         options('let a: A<B<C> & {} & {} & import("vue").ComponentCustomProperties) | null>;')
-      )
+      ).content
     ).toEqual(
       "import { ComponentCustomProperties } from 'vue';\nlet a: A<B<C> & {} & {} & ComponentCustomProperties) | null>;"
     )
@@ -59,13 +59,14 @@ describe('transform tests', () => {
         options(
           'declare const _default: import("./Service").ServiceConstructor<import("./Service").default>;'
         )
-      )
+      ).content
     ).toEqual(
       "import { ServiceConstructor, default as __DTS_DEFAULT_0__ } from './Service';\ndeclare const _default: ServiceConstructor<__DTS_DEFAULT_0__>;"
     )
 
     expect(
       transformCode(options('import { Type } from "./test";\nconst test: import("./test").Test;'))
+        .content
     ).toEqual("import { Type, Test } from './test';\n\nconst test: Test;")
   })
 
@@ -86,34 +87,34 @@ describe('transform tests', () => {
       clearPureImport: false
     })
 
-    expect(transformCode(options('import type { TestBase } from "@/src/test";'))).toEqual(
+    expect(transformCode(options('import type { TestBase } from "@/src/test";')).content).toEqual(
       "import { TestBase } from './test';\n"
     )
 
-    expect(transformCode(options('import("@/components/test").Test;'))).toEqual(
+    expect(transformCode(options('import("@/components/test").Test;')).content).toEqual(
       "import('../components/test').Test;"
     )
-    expect(transformCode(options('import type { TestBase } from "@/components/test";'))).toEqual(
-      "import { TestBase } from '../components/test';\n"
-    )
+    expect(
+      transformCode(options('import type { TestBase } from "@/components/test";')).content
+    ).toEqual("import { TestBase } from '../components/test';\n")
 
-    expect(transformCode(options('import("@/components/test").Test;\n'))).toEqual(
+    expect(transformCode(options('import("@/components/test").Test;\n')).content).toEqual(
       "import('../components/test').Test;\n"
     )
 
     expect(
       transformCode(
         options('import VContainer from "@components/layout/container/VContainer.vue";')
-      )
+      ).content
     ).toEqual(
       "import { default as VContainer } from './components/layout/container/VContainer.vue';\n"
     )
 
-    expect(transformCode(options('import type { TestBase } from "~/test";'))).toEqual(
+    expect(transformCode(options('import type { TestBase } from "~/test";')).content).toEqual(
       "import { TestBase } from './test';\n"
     )
 
-    expect(transformCode(options('import type { TestBase } from "$src/test";'))).toEqual(
+    expect(transformCode(options('import type { TestBase } from "$src/test";')).content).toEqual(
       "import { TestBase } from './test';\n"
     )
   })
@@ -128,12 +129,13 @@ describe('transform tests', () => {
       clearPureImport: true
     })
 
-    expect(transformCode(options('import "@/themes/common.scss";'))).toEqual('')
+    expect(transformCode(options('import "@/themes/common.scss";')).content).toEqual('')
     expect(
       transformCode(options('import "@/themes/common.scss";\nimport type { Ref } from "vue";'))
+        .content
     ).toEqual("import { Ref } from 'vue';\n")
     expect(
-      transformCode(options("{ 'database-import': import('vue').FunctionalComponent }"))
+      transformCode(options("{ 'database-import': import('vue').FunctionalComponent }")).content
     ).toEqual("{ 'database-import': import('vue').FunctionalComponent }")
   })