Skip to content
Hamidreza Ebtehaj edited this page Apr 14, 2023 · 8 revisions

Examples

Example from PR #49

Expand Example from PR #49
local NuiTree = require("nui.tree")
local Split = require("nui.split")
local NuiLine = require("nui.line")

local split = Split({
  relative = "win",
  position = "bottom",
  size = 30,
})

split:mount()

-- quit
split:map("n", "q", function()
  split:unmount()
end, { noremap = true })

local tree = NuiTree({
  winid = split.winid,
  nodes = {
    NuiTree.Node({ text = "a" }),
    NuiTree.Node({ text = "b" }, {
      NuiTree.Node({ text = "b-1" }),
      NuiTree.Node({ text = "b-2" }, {
        NuiTree.Node({ text = "b-1-a" }),
        NuiTree.Node({ text = "b-2-b" }),
      }),
    }),
    NuiTree.Node({ text = "c" }, {
      NuiTree.Node({ text = "c-1" }),
      NuiTree.Node({ text = "c-2" }),
    }),
  },
  prepare_node = function(node)
    local line = NuiLine()

    line:append(string.rep("  ", node:get_depth() - 1))

    if node:has_children() then
      line:append(node:is_expanded() and "" or "", "SpecialChar")
    else
      line:append("  ")
    end

    line:append(node.text)

    return line
  end,
})

local map_options = { noremap = true, nowait = true }

-- print current node
split:map("n", "<CR>", function()
  local node = tree:get_node()
  print(vim.inspect(node))
end, map_options)

-- collapse current node
split:map("n", "h", function()
  local node = tree:get_node()

  if node:collapse() then
    tree:render()
  end
end, map_options)

-- collapse all nodes
split:map("n", "H", function()
  local updated = false

  for _, node in pairs(tree.nodes.by_id) do
    updated = node:collapse() or updated
  end

  if updated then
    tree:render()
  end
end, map_options)

-- expand current node
split:map("n", "l", function()
  local node = tree:get_node()

  if node:expand() then
    tree:render()
  end
end, map_options)

-- expand all nodes
split:map("n", "L", function()
  local updated = false

  for _, node in pairs(tree.nodes.by_id) do
    updated = node:expand() or updated
  end

  if updated then
    tree:render()
  end
end, map_options)

-- add new node under current node
split:map("n", "a", function()
  local node = tree:get_node()
  tree:add_node(
    NuiTree.Node({ text = "d" }, {
      NuiTree.Node({ text = "d-1" }),
    }),
    node:get_id()
  )
  tree:render()
end, map_options)

-- delete current node
split:map("n", "d", function()
  local node = tree:get_node()
  tree:remove_node(node:get_id())
  tree:render()
end, map_options)

tree:render()

Treesitter Query Files Browser

open_treesitter_query_files_browser()
-- or
open_treesitter_query_files_browser("c", "go", "lua", "rust")
Expand Treesitter Query Files Browser
local Popup = require("nui.popup")
local Line = require("nui.line")
local Text = require("nui.text")
local Tree = require("nui.tree")
local event = require("nui.utils.autocmd").event

local function read_file(filename)
  local file, err = io.open(filename, "r")
  if not file then
    error(err)
  end
  local content = file:read("*a")
  io.close(file)
  return content
end

local function map(items, callback)
  local result = {}
  for i, item in ipairs(items) do
    result[i] = callback(item, i)
  end
  return result
end

local function for_each(items, callback)
  for i, item in ipairs(items) do
    callback(item, i)
  end
end

local mod = {}

local function list_query_sources(lang)
  local sources = {}
  local paths = vim.treesitter.get_query_files(lang, "*")

  local is_seen = {}

  for i, path in ipairs(paths) do
    if not is_seen[path] then
      is_seen[path] = true

      sources[i] = {
        lang = lang,
        filepath = path,
        plugin_name = vim.fn.fnamemodify(path, ":h:h:h:t"),
        query_name = vim.fn.fnamemodify(path, ":t:r"),
      }
    end
  end

  return sources
end

local function check_query_file_health(node)
  local query_content = read_file(node.filepath)

  local ok, err = pcall(vim.treesitter.query.parse_query, node.lang, query_content)

  if ok then
    node.ok = true
    node.err = nil
    node.err_position = nil
  else
    node.ok = false
    node.err = err
    node.err_position = string.match(node.err, "position (%d+)")
  end
end

local function get_query_source_nodes(langs)
  local nodes = {}

  for _, lang in ipairs(langs) do
    local query_sources = list_query_sources(lang)
    local children = map(query_sources, function(source)
      source.id = lang .. "-" .. source.filepath
      return Tree.Node(source)
    end)
    table.insert(nodes, Tree.Node({ text = lang }, children))
  end

  return nodes
end

function mod.open_treesitter_query_files_browser(...)
  local langs = { ... }

  if #langs == 0 then
    langs = require("nvim-treesitter.info").installed_parsers()
  end

  local popup = Popup({
    enter = true,
    position = "50%",
    size = {
      width = "80%",
      height = "60%",
    },
    border = {
      style = "rounded",
      text = {
        top = "Treesitter Query Files",
      },
    },
    buf_options = {
      readonly = true,
      modifiable = false,
    },
    win_options = {
      winhighlight = "Normal:Normal,FloatBorder:Normal",
    },
  })

  popup:mount()

  popup:on({ event.BufWinLeave }, function()
    vim.schedule(function()
      popup:unmount()
    end)
  end, { once = true })

  local tree = Tree({
    winid = popup.winid,
    nodes = get_query_source_nodes(langs),
    prepare_node = function(node)
      local line = Line()

      if node:has_children() then
        line:append(node:is_expanded() and "" or "", "SpecialChar")
        line:append("Language: " .. node.text)

        return line
      end

      line:append("  [")
      if node.ok then
        line:append(Text("", "DiagnosticSignInfo"))
      elseif node.err then
        line:append(Text("×", "DiagnosticSignError"))
      else
        line:append(" ")
      end
      line:append("] ")

      line:append(node.plugin_name .. " ::: " .. node.query_name, {
        virt_text = { { " ::: " .. vim.fn.fnamemodify(node.filepath, ":~"), "Folded" } },
        virt_text_pos = "eol",
      })

      return line
    end,
  })

  local map_options = { remap = false, nowait = true }

  -- exit
  popup:map("n", { "q", "<esc>" }, function()
    popup:unmount()
  end, map_options)

  -- collapse
  popup:map("n", "h", function()
    local node, linenr = tree:get_node()
    if not node:has_children() then
      node, linenr = tree:get_node(node:get_parent_id())
    end
    if node and node:collapse() then
      vim.api.nvim_win_set_cursor(popup.winid, { linenr, 0 })
      tree:render()
    end
  end, map_options)

  -- expand
  popup:map("n", "l", function()
    local node, linenr = tree:get_node()
    if not node:has_children() then
      node, linenr = tree:get_node(node:get_parent_id())
    end
    if node and node:expand() then
      if not node.checked then
        node.checked = true

        vim.schedule(function()
          for _, n in ipairs(tree:get_nodes(node:get_id())) do
            check_query_file_health(n)
          end
          tree:render()
        end)
      end

      vim.api.nvim_win_set_cursor(popup.winid, { linenr, 0 })
      tree:render()
    end
  end, map_options)

  -- open
  popup:map("n", "o", function()
    local node = tree:get_node()
    if not node.filepath then
      return
    end

    vim.cmd(
      string.format(
        "tab drop %s | %s",
        node.filepath,
        node.err_position and string.gsub(string.format([[goto %s]], node.err_position), " ", " ") or ""
      )
    )
  end, map_options)

  -- refresh
  popup:map("n", "r", function()
    local node = tree:get_node()
    vim.schedule(function()
      if node:has_children() then
        for_each(tree:get_nodes(node:get_id()), check_query_file_health)
      else
        check_query_file_health(node)
      end

      tree:render()
    end)
  end, map_options)

  tree:render()
end

return mod

Guides

How to get all expanded nodes?

local function get_expanded_nodes(tree)
  local nodes = {}

  local function process(node)
    if node:is_expanded() then
      table.insert(nodes, node)

      if node:has_children() then
        for _, node in tree:get_nodes(node:get_id()) do
          process(node)
        end
      end
    end
  end

  for _, node in tree:get_nodes() do
    process(node)
  end

  return nodes
end

How to collapse all nodes?

Building upon the prior example, you can then collapse all nodes with:

local function collapse_all_nodes(tree)
  local expanded = get_expanded_nodes(tree)
  for _, id in ipairs(expanded) do
    local node = tree:get_node(id)
    node:collapse(id)
  end
  -- If you want to expand the root
  -- local root = tree:get_nodes()[1]
  -- root:expand()
end