Vit 依赖收集插件

2026-01-293 min read前端技术
vue
text
// plugins/vite-plugin-collect-render.ts
import { createFilter } from '@rollup/pluginutils'
import { parse as parseAst } from '@typescript-eslint/parser'
import { TSESTree } from '@typescript-eslint/types'
import { writeFileSync, mkdirSync, readFileSync } from 'fs'
import { resolve, dirname } from 'path'

const PLUGIN_NAME = 'vite:collect-render'

export default function collectRenderPlugin(options: {
    include?: string | string[]
    exclude?: string | string[]
    output?: string
} = {}) {

    const filter = createFilter(
        options.include || 'src/**/*.{vue,ts,tsx}',
        options.exclude || 'node_modules'
    )

    const output = resolve(options.output || 'src/auto-collect-renders.ts')

    // 存储收集到的 { key, renderCode }
    const renders: Array<{ key: string; render: string; file: string }> = []

    /** 遍历 AST 收集 collectRender */
    function walkAst(node: TSESTree.Node, code: string, file: string) {
        if (
            node.type === 'CallExpression' &&
            node.callee.type === 'Identifier' &&
            node.callee.name === 'collectRender'
        ) {
            console.log('[collect-render] 收集文件 collectRender:', file)
            const args = node.arguments
            if (args.length === 2) {
                const keyArg = args[0]
                const renderArg = args[1]

                if (
                    keyArg.type === 'Literal' &&
                    typeof keyArg.value === 'string' &&
                    (renderArg.type === 'ArrowFunctionExpression' ||
                        renderArg.type === 'FunctionExpression')
                ) {
                    const key = keyArg.value
                    const [start, end] = renderArg.range!
                    const renderCode = code.slice(start, end)

                    // 去重：后面覆盖前面
                    const index = renders.findIndex(r => r.key === key)
                    if (index >= 0) renders[index] = { key, render: renderCode, file }
                    else renders.push({ key, render: renderCode, file })
                }
            }
        }

        // 递归 children
        for (const k in node) {
            const value = (node as any)[k]
            if (Array.isArray(value)) value.forEach(v => v && typeof v.type === 'string' && walkAst(v, code, file))
            else if (value && typeof value.type === 'string') walkAst(value, code, file)
        }
    }

    /** 收集单个文件的 collectRender */
    function collectFile(file: string, code?: string) {
        if (!filter(file)) return
        if (!code) code = readFileSync(file, 'utf-8')

        let scriptCode = code
        let isVue = false
        if (file.endsWith('.vue')) {
            isVue = true
            const scriptSetupMatch = code.match(/<script setup.*?>([sS]*?)</script>/i)
            const scriptMatch = code.match(/<script(?! setup).*?>([sS]*?)</script>/i)
            if (scriptSetupMatch) scriptCode = scriptSetupMatch[1]
            else if (scriptMatch) scriptCode = scriptMatch[1]
            else return
        }
        try {
            const ast = parseAst(scriptCode, {
                range: true,
                sourceType: 'module',
                ecmaVersion: 'latest',
                ecmaFeatures: { jsx: true }
            })
            walkAst(ast as any, scriptCode, file)
        } catch (e) {
            console.warn(`[collect-render] AST 解析失败: ${file}`, e)
        }
    }

    /** 生成 auto-collect-renders.ts 文件 */
    function generateFile() {
        if (renders.length === 0) return

        const exports = renders.map(r => `  '${r.key}': ${r.render}`).join(',
')
        const content = `
// 🚨 This file is auto-generated by vite-plugin-collect-render
// 🔄 Do not edit manually

export const collectedRenders = {
${exports}
} as const

export type CollectedRenderKeys = keyof typeof collectedRenders
    `.trim()

        mkdirSync(dirname(output), { recursive: true })
        writeFileSync(output, content, 'utf-8')
    }

    return {
        name: PLUGIN_NAME,
        enforce: 'pre' as const,

        transform(code: string, id: string) {
            collectFile(id, code)
            return null
        },

        buildEnd() {
            generateFile()
            this.info(`${PLUGIN_NAME}: Collected ${renders.length} render functions → ${output}`)
        },

        // ⚡ HMR: 文件修改时重新收集并触发更新
        handleHotUpdate({ file, server }) {
            if (!filter(file)) return
            collectFile(file)
            generateFile()

            const module = server.moduleGraph.getModuleById(output)
            if (module) {
                server.moduleGraph.invalidateModule(module)
                server.ws.send({
                    type: 'update',
                    updates: [
                        {
                            type: 'js-update',
                            path: '/' + output,
                            acceptedPath: '/' + output,
                            timestamp: Date.now(),
                        },
                    ],
                })
            }
            return []
        }
    }
}