LuaSnip中实现Python代码片段自动导入依赖的技术方案

在代码编辑过程中，自动处理依赖导入是一个能显著提升开发效率的功能。本文将详细介绍如何在LuaSnip中实现Python代码片段自动检查并添加必要导入语句的技术方案。

核心思路

该方案的核心是通过LuaSnip的pre_expand事件回调，在代码片段展开前执行检查逻辑。主要实现以下功能：

1. 检查当前缓冲区是否已存在指定导入语句
2. 若不存在则自动在文件头部添加
3. 正确处理光标位置和代码片段展开位置

关键技术点

1. 事件回调机制

LuaSnip提供了pre_expand事件，允许我们在代码片段展开前执行自定义逻辑。相比post_expand，pre_expand能更好地处理光标位置问题。

2. 缓冲区操作

使用Neovim API操作缓冲区内容：

• nvim_buf_get_lines获取当前缓冲区内容
• nvim_buf_set_text插入新的导入语句
• 注意避免全缓冲区重写，防止破坏LuaSnip的内部标记

3. 位置标记处理

关键难点在于正确处理代码片段的展开位置。当在文件头部插入内容时，需要使用extmark来调整原始展开位置：

if vim.deep_equal(args.expand_pos, {0,0}) then
    vim.api.nvim_buf_set_extmark(0, require("luasnip.session").ns_id, inserted_lines, 0, {id = args.expand_pos_mark_id})
end

完整实现方案

基础版本：单条导入语句

local function check_import(args, user_args)
    -- 参数检查
    if type(user_args[1]) ~= "string" then
        error("参数必须是字符串")
    end

    -- 获取缓冲区内容
    local lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
    local import_exists = false

    -- 检查是否已存在导入语句
    for _, line in ipairs(lines) do
        if line:match("^" .. user_args[1]) then
            import_exists = true
            break
        end
    end

    -- 不存在则插入
    if not import_exists then
        vim.api.nvim_buf_set_text(0, 0,0, 0,0, {user_args[1], ""})
        
        -- 处理位置标记
        if vim.deep_equal(args.expand_pos, {0,0}) then
            vim.api.nvim_buf_set_extmark(0, require("luasnip.session").ns_id, 1,0, {id = args.expand_pos_mark_id})
        end
    end
end

增强版本：支持多条导入语句

local function check_import(args, user_args)
    -- 参数检查
    if type(user_args) ~= "table" then
        error("参数必须是表")
    end

    local lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
    local inserted_lines = 0

    -- 遍历所有需要检查的导入语句
    for _, import_statement in ipairs(user_args) do
        if type(import_statement) ~= "string" then
            error("每个导入语句必须是字符串")
        end

        local import_exists = false

        -- 检查当前导入语句是否存在
        for _, line in ipairs(lines) do
            if line:match("^" .. import_statement) then
                import_exists = true
                break
            end
        end

        -- 不存在则插入
        if not import_exists then
            vim.api.nvim_buf_set_text(0, 0, 0, 0, 0, {import_statement, ""})
            inserted_lines = inserted_lines + 1
        end
    end

    -- 处理位置标记
    if vim.deep_equal(args.expand_pos, {0, 0}) then
        vim.api.nvim_buf_set_extmark(0, require("luasnip.session").ns_id, inserted_lines, 0, {id = args.expand_pos_mark_id})
    end
end

应用示例

基础使用

s({trig="plot_2d"},
    fmta([[
        fig, ax = plt.subplots(figsize=[12, 10])
        ax.plot(<>, <>)
        ax.grid()
        ax.set_xlabel("<>")
        ax.set_ylabel("<>")
        plt.show()
        <>
    ]], {
        i(1), i(2),
        i(3, "x-Axis"),
        i(4, "y-Axis"),
        i(0)
    }),
    {
        callbacks = {
            [-1] = {
                [events.pre_expand] = function(_, args)
                    check_import(args, {"import matplotlib.pyplot as plt"})
                end
            }
        }
    }
)

多条导入语句

s({trig="plot_3d_surface"},
    fmta([[
        # Create data
        x = np.linspace(<>, <>, 100)
        y = np.linspace(<>, <>, 100)
        X, Y = np.meshgrid(x, y)
        Z = <>

        # Plot
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.plot_surface(X, Y, Z, cmap='viridis')
        ax.set_xlabel('X Label')
        ax.set_ylabel('Y Label')
        ax.set_zlabel('Z Label')
        plt.show()
        <>
    ]], {
        i(1), i(2), i(3), i(4), i(5), i(0)
    }),
    {
        callbacks = {
            [-1] = {
                [events.pre_expand] = function(_, args)
                    check_import(args, {
                        "from mpl_toolkits.mplot3d import Axes3D", 
                        "import numpy as np"
                    })
                end
            }
        }
    }
)

注意事项

1. 正则匹配使用了^确保匹配行首，避免匹配到注释中的内容
2. 插入语句时添加了空行""保持代码整洁
3. 位置标记处理是关键，确保代码片段在正确位置展开
4. 函数参数检查增加了鲁棒性

总结

通过LuaSnip的事件回调机制，我们实现了Python代码片段自动检查并添加依赖导入的功能。该方案具有以下优点：

• 自动化程度高，减少手动操作
• 支持多条导入语句检查
• 正确处理光标位置
• 代码复用性强