nvim/lua/snippets/python.lua

dofile(core.snippets)

--- convert snake to pascal case
---@return string classname
local function py_file_name()
  local fn = file_name()
  local new = ""

  -- convert snake to pascal case
  for i = 1, #fn do
    if i == 1 then
      new = fn:sub(1, 1):upper()
    elseif fn:sub(i - 1, i - 1) == "_" then
      new = new..fn:sub(i, i):upper()
    elseif fn:sub(i, i) ~= "_" then
      new = new..fn:sub(i, i)
    end
  end

  return new
end

--- find the current class and return its name if not available defaults to
--- the file name
---@return string
local function python_class()
  -- set the starting node to the node under the cursor
  local node = require("nvim-treesitter.ts_utils").get_node_at_cursor()

  while node do
    -- check if we're in a class
    if node:type() == "class_definition" then
      -- find the class name in the class declaration and return it
      for i = 0, node:child_count() - 1 do
        local child = node:child(i)
        if child and child:type() == "identifier" then
          return vim.treesitter.get_node_text(child, 0)
        end
      end
    end
    node = node:parent()
  end

  -- if no class can be found default to the current file name
  return file_name()
end

return {
  s("class", {
    t("class "),
    c(1, {
      f(py_file_name, {}),
      i(0, "MyClass")
    }),
    c(2, {
      t(""),
      sn(nil, {
        t("("),
        i(1, "object"),
        t(")")
      })
    }),
    t({ ": ", "\t" }),
    i(0)
  }),

  s({ trig = [[fn\|def\|constr\|init]], trigEngine = "vim" }, {
    t("def "),
    d(1, function(_, snip)
      if snip.trigger == "constr" or snip.trigger == "init" then
        print(python_class())
        return sn(nil, {
          t("__init__")
        })
      else
        return sn(nil, {
          i(1, "myFunc")
        })
      end
    end, {}),
    t("("),
    d(2, function(_, snip)
      -- if this is a constructor or a function in a class then we need to put
      -- self as the first argument
      if snip.trigger == "constr" or snip.trigger == "init"
        or python_class() ~= file_name() then

        return sn(nil, {
          t("self")
        })
      else
        return sn(nil, {})
      end
    end, {}),
    i(3),
    t(")"),
    t(" -> "),
    i(4, "None"),
    t({ ":", "\t" }),
    i(0, "pass")
  }),

  s("main", {
    t({ "def main() -> None:", "\t" }),
    i(0, 'print("hello world")'),
    t({ "", "", 'if __name__ == "__main__":', "\tmain()" })
  })
}