如何实现Neovim命令自动生成Python函数参数打印语句？

一、基于Treesitter的深度解析实现 #

既然你已经能定位到函数参数节点，接下来可以通过遍历参数节点的子节点，提取不同类型的参数名称，再生成对应的打印语句。

核心步骤： #

1. 遍历参数节点子元素：针对Python函数的parameters节点，逐个处理子节点，区分普通参数、可变参数（*args）、关键字可变参数（**kwargs）以及解构参数。
2. 提取参数名称：
   • 普通参数：直接取identifier节点的文本内容
   • 可变参数：提取star节点后的identifier文本
   • 关键字可变参数：提取double_star节点后的identifier文本
   • 解构参数（如(a, b)或{x: x_val}）：递归解析内部变量名（可根据需求简化处理）
3. 生成打印语句：用f-string格式生成print(f"param: {param}")这类语句，批量插入到函数体的第一行位置。

示例Lua代码片段： #

local function get_param_names(node)
  local params = {}
  for child in node:iter_children() do
    local type = child:type()
    if type == "identifier" then
      table.insert(params, child:text())
    elseif type == "star" then
      local next_child = child:next_sibling()
      if next_child and next_child:type() == "identifier" then
        table.insert(params, "*" .. next_child:text())
      end
    elseif type == "double_star" then
      local next_child = child:next_sibling()
      if next_child and next_child:type() == "identifier" then
        table.insert(params, "**" .. next_child:text())
      end
    -- 可选：处理元组/列表解构参数
    elseif type == "tuple_pattern" or type == "list_pattern" then
      for inner_child in child:iter_children() do
        if inner_child:type() == "identifier" then
          table.insert(params, inner_child:text())
        end
      end
    end
  end
  return params
end

local function generate_print_statements(params)
  local lines = {}
  for _, param in ipairs(params) do
    table.insert(lines, string.format('print(f"%s: %r")', param, param))
  end
  return lines
end

-- 绑定到Neovim命令
vim.api.nvim_create_user_command("InsertParamPrints", function()
  local parser = vim.treesitter.get_parser(0, "python")
  local tree = parser:parse()[1]
  local root = tree:root()
  local cursor = vim.api.nvim_win_get_cursor(0)
  local row, col = cursor[1]-1, cursor[2]

  -- 定位光标所在的函数节点
  local function_node = nil
  local function_query = vim.treesitter.query.parse("python", [[
    (function_definition) @func
  ]])
  for _, match, _ in query:iter_matches(root, 0) do
    local node = match[1]
    local start_row, start_col, end_row, end_col = node:range()
    if (row > start_row or (row == start_row and col >= start_col)) and
       (row < end_row or (row == end_row and col <= end_col)) then
      function_node = node
      break
    end
  end

  if not function_node then
    vim.notify("Not inside a Python function")
    return
  end

  -- 获取参数节点并生成打印语句
  local params_node = function_node:field("parameters")[1]
  local params = get_param_names(params_node)
  local print_lines = generate_print_statements(params)

  -- 插入到函数体第一行
  local body_node = function_node:field("body")[1]
  local insert_row = body_node:range() + 1
  vim.api.nvim_buf_set_lines(0, insert_row, insert_row, false, print_lines)
end, {})

二、用LSP API简化参数识别 #

如果你的Neovim已经配置了Python LSP（如pyright、pylsp），可以直接通过LSP获取函数参数信息，无需手动解析语法树，更简便且准确。

核心思路： #

1. 发送textDocument/documentSymbol请求：获取当前文档的所有符号，找到光标所在的函数符号。
2. 提取参数列表：从函数符号的parameters字段中直接获取参数名称（LSP会自动处理类型注解、默认值等情况）。
3. 生成并插入打印语句：和Treesitter方案逻辑一致，生成打印语句插入到函数体开头。

示例Lua代码片段： #

vim.api.nvim_create_user_command("InsertParamPrintsLSP", function()
  local bufnr = vim.api.nvim_get_current_buf()
  local cursor = vim.api.nvim_win_get_cursor(0)
  local position = {line = cursor[1]-1, character = cursor[2]}

  -- 发送LSP请求获取文档符号
  vim.lsp.buf.request(bufnr, "textDocument/documentSymbol", {
    textDocument = vim.lsp.util.make_text_document_params(bufnr)
  }, function(err, result)
    if err then
      vim.notify("LSP request failed: " .. err.message)
      return
    end

    -- 递归查找光标所在的函数符号
    local function find_function_symbol(symbols)
      for _, sym in ipairs(symbols) do
        local range = sym.range
        local start = range.start
        local end_ = range["end"]
        local in_range = (position.line > start.line or (position.line == start.line and position.character >= start.character)) and
                        (position.line < end_.line or (position.line == end_.line and position.character <= end_.character))
        if in_range then
          if sym.kind == 12 then -- 12对应LSP定义的Function类型
            return sym
          elseif sym.children then
            local found = find_function_symbol(sym.children)
            if found then return found end
          end
        end
      end
      return nil
    end

    local func_sym = find_function_symbol(result)
    if not func_sym then
      vim.notify("Not inside a Python function")
      return
    end

    -- 提取参数并生成打印语句
    local params = {}
    if func_sym.parameters then
      for _, param in ipairs(func_sym.parameters) do
        table.insert(params, param.name)
      end
    end

    local print_lines = {}
    for _, param in ipairs(params) do
      table.insert(print_lines, string.format('print(f"%s: %r")', param, param))
    end

    -- 插入到函数体第一行
    local insert_row = func_sym.range["end"].line + 1
    vim.api.nvim_buf_set_lines(bufnr, insert_row, insert_row, false, print_lines)
  end)
end, {})

两种方案对比 #

• Treesitter方案：无需依赖LSP服务，离线可用，但需要手动处理各种Python参数语法细节（如解构、类型注解）。
• LSP方案：参数识别更准确，自动处理复杂语法，代码逻辑更简洁，但必须确保Python LSP服务正常运行。

内容的提问来源于stack exchange，提问作者Razumov