Automatically add async to function declaration when typing "await"

typescript.lua

---@param types string[] Will return the first node that matches one of these types
---@param node TSNode|nil
---@return TSNode|nil
local function find_node_ancestor(types, node)
  if not node then
    return nil
  end

  if vim.tbl_contains(types, node:type()) then
    return node
  end

  local parent = node:parent()

  return find_node_ancestor(types, parent)
end


---When typing "await" add "async" to the function declaration if the function
---isn't async already.
local function add_async()
  -- This function should be executed when the user types "t" in insert mode,
  -- but "t" is not inserted because it's the trigger.
  vim.api.nvim_feedkeys('t', 'n', true)

  local buffer = vim.fn.bufnr()

  local text_before_cursor = vim.fn.getline('.'):sub(vim.fn.col '.' - 4, vim.fn.col '.' - 1)
  if text_before_cursor ~= 'awai' then
    return
  end

  -- ignore_injections = false makes this snippet work in filetypes where JS is injected
  -- into other languages
  local current_node = vim.treesitter.get_node { ignore_injections = false }
  local function_node = find_node_ancestor(
    { 'arrow_function', 'function_declaration', 'function', 'method_definition' },
    current_node
  )
  if not function_node then
    return
  end

  local function_text = vim.treesitter.get_node_text(function_node, 0)
  if vim.startswith(function_text, 'async') then
    return
  end

  local start_row, start_col = function_node:start()
  vim.api.nvim_buf_set_text(buffer, start_row, start_col, start_row, start_col, { 'async ' })
end

vim.keymap.set('i', 't', add_async, { buffer = true })

Great plugin.
Additionally, you could include the method_definition ancestor so it will work also with classes:

 local function_node = find_node_ancestor(
    { 'arrow_function', 'function_declaration', 'function', 'method_definition' },
    current_node
  )

Great plugin. Additionally, you could include the method_definition ancestor so it will work also with classes:

 local function_node = find_node_ancestor(
    { 'arrow_function', 'function_declaration', 'function', 'method_definition' },
    current_node
  )

Looks great, thanks! Updated the Gist 🎉

Hi @JoosepAlviste, thanks for the gist! Is there an easy way to identify if the cursor is currently in a comment when triggering the keymap t in insert mode? I'm not having any luck with my attempts, such as

-- 1: doesn't work, nodetype is `statement_block`
vim.keymap.set("i", "t", function()
  if vim.treesitter.get_node():type() == "comment" then
    return
  end

  add_async()
end, { buffer = true, desc = "trigger add_async" })

-- 2: also doesn't work, attempting to reparse the unknown tree before getting node
vim.keymap.set("i", "t", function()
  vim.treesitter.get_parser():parse()
  if vim.treesitter.get_node():type() == "comment" then
    return
  end

  add_async()
end, { buffer = true, desc = "trigger add_async" })

For example, leaving a comment mentioning "await". Typing within this comment will unnecessarily add async to bar()

class FooClass {
  bar() {
    // we don't need to await[CURSOR HERE]
    const baz = doThing();
  }
}

Hey! I made another addition for the sake of completion of this amazing script. I have modified the script so that it also works for access modified functions inside of classes. I was using this script all the time at work, and since most of the functions I write are inside classes, it made sense to me to modify it for the public, private and protected modifiers. Here is the updated script:

---@param types string[] Will return the first node that matches one of these types
---@param node TSNode|nil
---@return TSNode|nil
local function find_node_ancestor(types, node)
  if not node then
    return nil
  end

  if vim.tbl_contains(types, node:type()) then
    return node
  end

  local parent = node:parent()

  return find_node_ancestor(types, parent)
end


---When typing "await" add "async" to the function declaration if the function
---isn't async already.
local function add_async()
  -- This function should be executed when the user types "t" in insert mode,
  -- but "t" is not inserted because it's the trigger.
  vim.api.nvim_feedkeys('t', 'n', true)

  local buffer = vim.fn.bufnr()

  local text_before_cursor = vim.fn.getline('.'):sub(vim.fn.col '.' - 4, vim.fn.col '.' - 1)
  if text_before_cursor ~= 'awai' then
    return
  end

  -- ignore_injections = false makes this snippet work in filetypes where JS is injected
  -- into other languages
  local current_node = vim.treesitter.get_node { ignore_injections = false }
  local function_node = find_node_ancestor(
    { 'arrow_function', 'function_declaration', 'function', 'method_definition' },
    current_node
  )
  if not function_node then
    return
  end

  local function_text = vim.treesitter.get_node_text(function_node, 0)
  if vim.startswith(function_text, 'async') then
    return
  end

  if vim.startswith(function_text, 'public async') then
    return
  end
  if vim.startswith(function_text, 'private async') then
    return
  end
  if vim.startswith(function_text, 'protected async') then
    return
  end

  local start_row, start_col = function_node:start()

  if vim.startswith(function_text, 'public') then
    start_col = start_col + 7
  end
  if vim.startswith(function_text, 'private') then
    start_col = start_col + 8
  end
  if vim.startswith(function_text, 'protected') then
    start_col = start_col + 10
  end

  vim.api.nvim_buf_set_text(buffer, start_row, start_col, start_row, start_col, { 'async ' })
end

vim.keymap.set('i', 't', add_async, { buffer = true })