Neovim: Toggle f-string for python using treesitter

toggle_fstring.lua

local ts_utils = require("nvim-treesitter.ts_utils")

local M = {}

-- Credit: https://www.reddit.com/r/neovim/comments/tge2ty/python_toggle_fstring_using_treesitter/
M.toggle_fstring = function()
  local winnr = 0
  local cursor = vim.api.nvim_win_get_cursor(winnr)
  local node = ts_utils.get_node_at_cursor()  ---@type TSNode?

  while (node ~= nil) and (node:type() ~= "string") do
    node = node:parent()
  end
  if node == nil then
    vim.cmd.echon [["f-string: not in a string node."]]
    return
  end

  ---@diagnostic disable-next-line: unused-local
  local srow, scol, ecol, erow = ts_utils.get_vim_range({ node:range() })
  vim.fn.setcursorcharpos(srow, scol)

  local char = vim.api.nvim_get_current_line():sub(scol, scol)
  local is_fstring = (char == "f")

  if is_fstring then
    vim.cmd [[normal "_x]]
    -- if cursor is in the same line as text change
    if srow == cursor[1] then
      cursor[2] = cursor[2] - 1 -- negative offset to cursor
    end
  else
    vim.cmd [[noautocmd normal if]]
    -- if cursor is in the same line as text change
    if srow == cursor[1] then
      cursor[2] = cursor[2] + 1 -- positive offset to cursor
    end
  end
  vim.api.nvim_win_set_cursor(winnr, cursor)
end

return M