re-plugin vite auto-import / vue-components

extract-components.ts

import { type SFCTemplateBlock } from '@vue/compiler-sfc';
import { NodeTypes, Node, ElementNode, ElementTypes } from '@vue/compiler-core';

// 提取 vue 文件中的组件
export function extractComponents(
    tpl: SFCTemplateBlock,
    filter: (name: string) => boolean,
) {
    const list: ElementNode[] = [];
    if (!tpl) {
        return list;
    }
    const root = tpl.ast;
    const visit = (node: Node) => {
        if (node.type === NodeTypes.ELEMENT) {
            const element = node as ElementNode;
            if (
                element.tagType === ElementTypes.COMPONENT &&
                filter(element.tag)
            ) {
                list.push(element);
            }
        }
        if ('children' in node) {
            const children = node.children as Node[];
            for (const n of children) {
                visit(n);
            }
        }
    };

    visit(root);

    return list;
}

extract-globals.ts

// 提取 ts 中的 import，以及使用到了的全局变量
// TODO 考虑不做，直接使用 vscode 的自动引入，主要就是 vue 的常用方法
import { parse, type SFCScriptBlock } from '@vue/compiler-sfc';
import * as ts from 'typescript';
export function extractGlobals(script: SFCScriptBlock) {
    const importNames = new Set<string>();
    if (!script) {
        return importNames;
    }
    const { statements } = ts.createSourceFile(
        'index.d.ts',
        script.content,
        ts.ScriptTarget.ESNext,
        true,
        ts.ScriptKind.TS,
    );

    // TODO 用户自定义的变量就先不考虑？
    for (const s of statements) {
        // import { useI18n } from 'vue-i18n';
        // import VueI18n from 'vue-i18n';
        if (!ts.isImportDeclaration(s)) {
            continue;
        }
        // { useI18n } or VueI18n
        const { importClause } = s;
        if (!importClause) {
            continue;
        }
        const {
            // VueI18n
            name,
            // { useI18n }
            namedBindings,
        } = importClause;

        if (name) {
            importNames.add(name.text);
        }
        if (namedBindings) {
            if (ts.isNamedImports(namedBindings)) {
                for (const n of namedBindings.elements) {
                    importNames.add(n.name.text);
                }
            } else if (ts.isNamespaceImport(namedBindings)) {
                // import * as Name
                importNames.add(namedBindings.name.text);
            }
        }
    }

    return importNames;
}

index.ts

/**
 * 尝试批量处理 unplugin 代码，将它恢复为显式 import 写法
 *
 * 目前暂只支持解析 components unplugin
 *
 * 用法:
 * ```
 * npx tsx scripts/re-plugin packages/sdk-chart/src/
 * ```
 *
 * 会将目录下的所有 .vue 文件都进行一次处理
 *
 * 为什么不用 unplugin?
 *   - 隐式声明使得代码不够明确
 *   - 当项目变大时，unplugin 冲突问题明显，造成更大的隐患
 *
 * 包含的内容
 *   - vue 常用方法，例如 ref、computed 等
 *   - antdv 组件，例如 a-input
 *   - 业务组件，例如 color-picker
 */

import { parse } from '@vue/compiler-sfc';
import { readdir, readFile, stat, writeFile } from 'node:fs/promises';
import { join } from 'node:path';
import { readUnpluginComponents, Source } from './read-unplugin';
import { pascalCase } from './utils';
import assert from 'node:assert';
import { extractComponents } from './extract-components';
import { extractGlobals } from './extract-globals';

const UNPLUGIN_SRC_DIR = join(__dirname, '../../src');

interface Import extends Source {
    name: string;
}

async function main() {
    const components = await readUnpluginComponents(UNPLUGIN_SRC_DIR);
    const dir = process.argv[2];
    if (!dir) {
        throw new Error(
            `请指定要解析的 .vue 文件所在目录，例如 npx tsx scripts/re-plugin packages/sdk-chart/src`,
        );
    }

    const walk = async (dir: string) => {
        const list = await readdir(dir);

        for (const file of list) {
            const filePath = join(dir, file);
            const s = await stat(filePath);
            if (s.isDirectory()) {
                await walk(filePath);
            } else if (file.endsWith('.vue')) {
                rePluginVueFile(filePath, components);
            }
        }
    };

    await walk(join(__dirname, '../../', dir));
}

async function rePluginVueFile(
    filePath: string,
    components: Awaited<ReturnType<typeof readUnpluginComponents>>,
) {
    const content = await readFile(filePath, 'utf-8');
    const result = parse(content);
    // console.dir(result.descriptor.template.ast)

    const { template, script, scriptSetup } = result.descriptor;

    const list = extractComponents(result.descriptor.template, (n) =>
        components.has(n),
    );

    const globals = new Set([
        ...extractGlobals(script),
        ...extractGlobals(scriptSetup),
    ]);

    const imports = generateImportStatements(
        list
            .map<Import>((n) => {
                const source = components.get(n.tag)!;
                return {
                    ...source,
                    name: pascalCase(n.tag),
                };
            })
            .filter((o) => !globals.has(o.name)),
    );

    const importStatements = imports.join(';\n');

    // 回写文件
    if (!importStatements) {
        console.log(`无需修改 - ${filePath}`);
        return;
    }
    const { source } = result.descriptor;
    const { loc } = scriptSetup ?? script ?? {};
    await writeFile(
        filePath,
        source.slice(0, loc.start.offset) +
            '\n' +
            importStatements +
            source.slice(loc.start.offset),
        'utf8',
    );
    console.log(
        `已更新 - ${filePath}\n${imports.map((l) => `  ${l}`).join('\n')}`,
    );
}

// 生成 import 语句
function generateImportStatements(list: Import[]) {
    // TODO 未来可以解析已经导入过了的
    const imports = new Map<string, Map<string, string>>();

    for (const { module, path, name } of list) {
        let members = imports.get(module);
        if (!members) {
            members = new Map<string, string>();
            imports.set(module, members);
        }
        let member = members.get(name);
        if (member) {
            assert(member === path);
            continue;
        }
        members.set(name, path);
    }

    const lines: string[] = [];
    for (const [m, members] of imports) {
        const DEFAULT = 'default';

        let defaultStatement = '';
        const others: string[] = [];
        for (const [name, path] of members) {
            if (path === DEFAULT) {
                defaultStatement = `${name}`;
                continue;
            }
            others.push(name === path ? name : `${path} as ${name}`);
        }
        const otherStatement = others.length ? `{ ${others.join(', ')} }` : '';
        lines.push(
            `import ${[defaultStatement, otherStatement]
                .filter(Boolean)
                .join(', ')} from '${m}'`,
        );
    }
    return lines;
}

main().catch(console.error);

read-unplugin.ts

// 解析 src 下的 auto-import.d.ts、components.d.ts，获取映射列表

import assert from 'node:assert';
import { readFile } from 'node:fs/promises';
import { join } from 'node:path';
import * as ts from 'typescript';
import { pascalCase } from './utils';
import { kebabCase } from 'lodash-es';

export interface Source {
    /** 从哪个 module 引入，例如 ant-design-vue、./components/ColorPicker.vue */
    module: string;
    /**
     * 导出的访问路径，例如 default、 ColorPicker
     *
     * 暂时我们只处理一层的场景吧
     */
    path: string;
}

export async function readUnpluginComponents(srcDir: string) {
    const filePath = join(srcDir, 'components.d.ts');
    const content = await readFile(filePath, 'utf-8');

    const { statements } = ts.createSourceFile(
        'index.d.ts',
        content,
        ts.ScriptTarget.ESNext,
        true,
        ts.ScriptKind.TS,
    );

    const parsed = new Map<string, Source>();

    for (const s of statements) {
        if (!ts.isModuleDeclaration(s)) {
            continue;
        }

        const { body } = s;
        // declare module 'vue' { ... }
        assert(ts.isModuleBlock(body));
        const {
            statements: [interfaceDeclaration, ...rest],
        } = body;

        // export interface GlobalComponents { ... }
        assert(
            ts.isInterfaceDeclaration(interfaceDeclaration) &&
                interfaceDeclaration.name.text === 'GlobalComponents',
        );
        assert(!rest.length, '不应该有多余语句了');

        const { members } = interfaceDeclaration;

        for (const m of members) {
            // AAvatar: typeof import('ant-design-vue/es')['Avatar']
            assert(ts.isPropertySignature(m));
            // AAvatar
            const name = m.name.getText();
            const { type } = m;
            // typeof import('ant-design-vue/es')['Avatar']
            assert(ts.isIndexedAccessTypeNode(type));
            const { objectType, indexType } = type;
            // typeof import('ant-design-vue/es')
            assert(ts.isImportTypeNode(objectType) && objectType.isTypeOf);
            const { argument } = objectType;
            assert(ts.isLiteralTypeNode(argument));
            const module = argument.literal;
            assert(ts.isStringLiteral(module));

            // ['Avatar']
            assert(ts.isLiteralTypeNode(indexType));
            const index = indexType.literal;
            assert(ts.isStringLiteral(index));

            const source: Source = {
                module: fixModule(module.text),
                path: index.text,
            };

            parsed.set(pascalCase(name), source);
            parsed.set(kebabCase(name), source);
        }
    }
    return parsed;
}

function fixModule(module: string) {
    if (module === 'ant-design-vue/es') {
        return 'ant-design-vue';
    }
    return module;
}

utils.ts

import { camelCase } from 'lodash-es';

export function pascalCase(str: string) {
    return capitalize(camelCase(str));
}

function capitalize(str: string) {
    return str.charAt(0).toUpperCase() + str.slice(1);
}