135 lines
2.9 KiB
Lua
135 lines
2.9 KiB
Lua
local ls = require "luasnip"
|
|
|
|
local snippet_from_nodes = ls.sn
|
|
local c = ls.choice_node
|
|
local t = ls.text_node
|
|
local d = ls.dynamic_node
|
|
local i = ls.insert_node
|
|
local rep = require("luasnip.extras").rep
|
|
local fmt = require("luasnip.extras.fmt").fmt
|
|
|
|
local ts_locals = require "nvim-treesitter.locals"
|
|
local ts_utils = require "nvim-treesitter.ts_utils"
|
|
|
|
local get_node_text = vim.treesitter.get_node_text
|
|
|
|
vim.treesitter.query.set(
|
|
"go",
|
|
"LuaSnip_Result",
|
|
[[
|
|
[
|
|
(method_declaration result: (_) @id)
|
|
(function_declaration result: (_) @id)
|
|
(func_literal result: (_) @id)
|
|
]
|
|
]]
|
|
)
|
|
|
|
local transform = function(text, info)
|
|
if text == "int" then
|
|
return t "0"
|
|
elseif text == "error" then
|
|
if info then
|
|
info.index = info.index + 1
|
|
|
|
return c(info.index, {
|
|
t(string.format('fmt.Errorf("%s %%w", %s)', info.func_name, info.err_name)),
|
|
t(info.err_name),
|
|
})
|
|
else
|
|
return t "err"
|
|
end
|
|
elseif text == "bool" then
|
|
return t "false"
|
|
elseif text == "string" then
|
|
return t '""'
|
|
elseif string.find(text, "*", 1, true) then
|
|
return t "nil"
|
|
end
|
|
|
|
return t(text)
|
|
end
|
|
|
|
local handlers = {
|
|
["parameter_list"] = function(node, info)
|
|
local result = {}
|
|
|
|
local count = node:named_child_count()
|
|
for idx = 0, count - 1 do
|
|
table.insert(result, transform(get_node_text(node:named_child(idx), 0), info))
|
|
if idx ~= count - 1 then
|
|
table.insert(result, t { ", " })
|
|
end
|
|
end
|
|
|
|
return result
|
|
end,
|
|
|
|
["type_identifier"] = function(node, info)
|
|
local text = get_node_text(node, 0)
|
|
return { transform(text, info) }
|
|
end,
|
|
}
|
|
|
|
local function go_result_type(info)
|
|
local cursor_node = ts_utils.get_node_at_cursor()
|
|
local scope = ts_locals.get_scope_tree(cursor_node, 0)
|
|
|
|
local function_node
|
|
for _, v in ipairs(scope) do
|
|
if v:type() == "function_declaration" or v:type() == "method_declaration" or v:type() == "func_literal" then
|
|
function_node = v
|
|
break
|
|
end
|
|
end
|
|
|
|
local query = vim.treesitter.query.get("go", "LuaSnip_Result")
|
|
for _, node in query:iter_captures(function_node, 0) do
|
|
if handlers[node:type()] then
|
|
return handlers[node:type()](node, info)
|
|
end
|
|
end
|
|
end
|
|
|
|
local go_ret_vals = function(args)
|
|
local info = { index = 0, err_name = args[1][1], func_name = args[2][1] }
|
|
return snippet_from_nodes(nil, go_result_type(info))
|
|
end
|
|
|
|
local M = {
|
|
main = {
|
|
t { "func main() {", "\t" },
|
|
i(0),
|
|
t { "", "}" },
|
|
},
|
|
ef = {
|
|
i(1, { "val" }),
|
|
t ", err := ",
|
|
i(2, { "f" }),
|
|
t "(",
|
|
i(3),
|
|
t ")",
|
|
i(0),
|
|
},
|
|
efi = {
|
|
i(1, { "val" }),
|
|
", ",
|
|
i(2, { "err" }),
|
|
" := ",
|
|
i(3, { "f" }),
|
|
"(",
|
|
i(4),
|
|
")",
|
|
t { "", "if " },
|
|
rep(2),
|
|
t { " != nil {", "\treturn " },
|
|
d(5, go_ret_vals, { 2, 3 }),
|
|
t { "", "}" },
|
|
i(0),
|
|
},
|
|
}
|
|
|
|
M.f = fmt("func {}({}) {} {{\n\t{}\n}}", { i(1, "name"), i(2), i(3), i(0) })
|
|
|
|
return M
|