Refactor: modernize plugin, deduplicate code, improve robustness

lua/pytrize/api.lua

@@ -1,21 +1,18 @@
 local M = {}
 
+--- Clear pytrize virtual text from buffer.
+---@param bufnr? integer Buffer number (0 or nil for current buffer)
 M.clear = function(bufnr)
     local marks = require("pytrize.marks")
-
-    if bufnr == nil then
-        bufnr = 0
-    end
-    marks.clear(bufnr)
+    marks.clear(bufnr or 0)
 end
 
+--- Set pytrize virtual text for parametrize entries in buffer.
+---@param bufnr? integer Buffer number (0 or nil for current buffer)
 M.set = function(bufnr)
     local cs = require("pytrize.call_spec")
     local marks = require("pytrize.marks")
-
-    if bufnr == nil then
-        bufnr = 0
-    end
+    bufnr = bufnr or 0
     marks.clear(bufnr)
     local call_specs_per_func = cs.get_calls(bufnr)
     if call_specs_per_func == nil then
@@ -45,24 +42,28 @@ M.set = function(bufnr)
     end
 end
 
+--- Jump to the parametrize entry declaration under cursor.
 M.jump = function()
     local jump = require("pytrize.jump")
 
     jump.to_param_declaration()
 end
 
+--- Jump to fixture definition under cursor.
 M.jump_fixture = function()
     local jump = require("pytrize.jump")
 
     jump.to_fixture_declaration()
 end
 
+--- Rename fixture under cursor across the project.
 M.rename_fixture = function()
     local rename = require("pytrize.rename")
 
     rename.rename_fixture()
 end
 
+--- Show all usages of fixture under cursor in quickfix list.
 M.fixture_usages = function()
     local usages = require("pytrize.usages")
 

lua/pytrize/call_spec.lua

@@ -1,34 +1,44 @@
 local M = {}
 
 local ts = vim.treesitter
-local ts_query = ts.query
-local parse_query = ts_query.parse or ts_query.parse_query
+local parse_query = vim.treesitter.query.parse
 
 local warn = require("pytrize.warn").warn
 local tbls = require("pytrize.tables")
 
 local get_root = function(bufnr)
-    local parser = ts.get_parser(bufnr)
-    local tstree = parser:parse()[1]
-    return tstree:root()
+    local ok, parser = pcall(ts.get_parser, bufnr, "python")
+    if not ok then
+        warn("Python treesitter parser not available")
+        return nil
+    end
+    local trees = parser:parse()
+    if not trees or #trees == 0 then
+        warn("Failed to parse buffer")
+        return nil
+    end
+    return trees[1]:root()
 end
 
+local PARAM_QUERY = parse_query(
+    "python",
+    [[
+      (decorated_definition (
+        decorator (
+          call
+            function: ((attribute) @param)
+        )
+      ))
+    ]]
+)
+
 local get_param_call_nodes = function(bufnr)
     local tsroot = get_root(bufnr)
-    local query = parse_query(
-        "python",
-        -- TODO not sure why eg (#eq? @param "pytest.mark.parametrize") does not work
-        [[
-          (decorated_definition (
-            decorator (
-              call
-                function: ((attribute) @param)
-            )
-          ))
-        ]]
-    )
+    if tsroot == nil then
+        return {}
+    end
     local nodes = {}
-    for _, node, _ in query:iter_captures(tsroot) do
+    for _, node, _ in PARAM_QUERY:iter_captures(tsroot, bufnr) do
         if ts.get_node_text(node, bufnr) == "pytest.mark.parametrize" then
             table.insert(nodes, node:parent())
         end
@@ -38,7 +48,17 @@ end
 
 local get_second_arg_node = function(call_node)
     local arguments = call_node:field("arguments")[1]
-    return arguments:child(3)
+    if not arguments then
+        return nil
+    end
+    -- Use named children to avoid depending on raw child indices
+    local named = {}
+    for child in arguments:iter_children() do
+        if child:named() then
+            table.insert(named, child)
+        end
+    end
+    return named[2]
 end
 
 local get_named_children = function(node)
@@ -53,7 +73,7 @@ end
 
 local list_entries = function(call_node)
     local list = get_second_arg_node(call_node)
-    if list:type() ~= "list" then
+    if not list or list:type() ~= "list" then
         return {}
     end
     return get_named_children(list)
@@ -166,14 +186,21 @@ M.get_calls = function(bufnr)
         local func_name = ts.get_node_text(func:field("name")[1], bufnr)
 
         local arguments = call:field("arguments")[1]
-        local params_node = arguments:child(1)
-        if params_node:type() ~= "string" then
+        local first_named = nil
+        for child in arguments:iter_children() do
+            if child:named() then
+                first_named = child
+                break
+            end
+        end
+        local params_node = first_named
+        if not params_node or params_node:type() ~= "string" then
             local row = call:start()
             warn(
                 string.format(
                     "couldn't parse params (line %d)\n  expected `string`\n  got `%s`",
                     row,
-                    params_node:type()
+                    params_node and params_node:type() or "nil"
                 )
             )
             return

lua/pytrize/health.lua

@@ -0,0 +1,39 @@
+local M = {}
+
+M.check = function()
+    vim.health.start("pytrize")
+
+    -- Check Neovim version
+    if vim.fn.has("nvim-0.9.0") == 1 then
+        vim.health.ok("Neovim >= 0.9.0")
+    else
+        vim.health.error("Neovim >= 0.9.0 required")
+    end
+
+    -- Check Python treesitter parser
+    local ok = pcall(vim.treesitter.language.inspect, "python")
+    if ok then
+        vim.health.ok("Python treesitter parser installed")
+    else
+        vim.health.error("Python treesitter parser not found", {
+            "Install with :TSInstall python (nvim-treesitter) or compile manually",
+        })
+    end
+
+    -- Check for grep (used by fixture rename/usages)
+    if vim.fn.executable("grep") == 1 then
+        vim.health.ok("grep command available")
+    else
+        vim.health.warn("grep not found — fixture rename and usages features will not work")
+    end
+
+    -- Check highlight group
+    local settings = require("pytrize.settings").settings
+    if vim.fn.hlexists(settings.highlight) == 1 then
+        vim.health.ok(string.format("Highlight group '%s' exists", settings.highlight))
+    else
+        vim.health.warn(string.format("Highlight group '%s' not found — virtual text may be invisible", settings.highlight))
+    end
+end
+
+return M

lua/pytrize/init.lua

@@ -3,18 +3,30 @@ local M = {}
 local settings = require("pytrize.settings")
 
 local function setup_commands()
-    vim.cmd('command Pytrize lua require("pytrize.api").set()')
-    vim.cmd('command PytrizeClear lua require("pytrize.api").clear()')
-    vim.cmd('command PytrizeJump lua require("pytrize.api").jump()')
-    vim.cmd('command PytrizeJumpFixture lua require("pytrize.api").jump_fixture()')
-    vim.cmd('command PytrizeRenameFixture lua require("pytrize.api").rename_fixture()')
-    vim.cmd('command PytrizeFixtureUsages lua require("pytrize.api").fixture_usages()')
+    vim.api.nvim_create_user_command("Pytrize", function()
+        require("pytrize.api").set()
+    end, { desc = "Set pytrize virtual text for parametrize entries" })
+    vim.api.nvim_create_user_command("PytrizeClear", function()
+        require("pytrize.api").clear()
+    end, { desc = "Clear pytrize virtual text" })
+    vim.api.nvim_create_user_command("PytrizeJump", function()
+        require("pytrize.api").jump()
+    end, { desc = "Jump to parametrize entry under cursor" })
+    vim.api.nvim_create_user_command("PytrizeJumpFixture", function()
+        require("pytrize.api").jump_fixture()
+    end, { desc = "Jump to fixture definition under cursor" })
+    vim.api.nvim_create_user_command("PytrizeRenameFixture", function()
+        require("pytrize.api").rename_fixture()
+    end, { desc = "Rename fixture under cursor across project" })
+    vim.api.nvim_create_user_command("PytrizeFixtureUsages", function()
+        require("pytrize.api").fixture_usages()
+    end, { desc = "Show fixture usages in quickfix list" })
 end
 
+--- Configure the pytrize plugin.
+---@param opts? PytrizeSettings
 M.setup = function(opts)
-    if opts == nil then
-        opts = {}
-    end
+    opts = opts or {}
     settings.update(opts)
     if not settings.settings.no_commands then
         setup_commands()

lua/pytrize/jump/param.lua

@@ -7,8 +7,6 @@ local paths = require('pytrize.paths')
 local warn = require('pytrize.warn').warn
 local open_file = require('pytrize.jump.util').open_file
 local get_nodeids_path = require('pytrize.paths').get_nodeids_path
-local min = require('pytrize.utils').min
-local max = require('pytrize.utils').max
 
 local function query_file(func_name, callback)
   local rootdir, _ = paths.split_at_root(vim.api.nvim_buf_get_name(0))
@@ -55,7 +53,7 @@ local function jump_to_nodeid_at_cursor(callback)
   local nodeid = nids.parse_raw(line:sub(i))
   local pattern_position = col_num + 1 - (i - 1)  -- cursor relative to match
   local param_position = pattern_position - nodeid.param_start_idx + 1  -- cursor relative to params
-  param_position = min(max(1, param_position), nodeid.params:len())  -- restrict it to be inside
+  param_position = math.min(math.max(1, param_position), nodeid.params:len())  -- restrict it to be inside
   if nodeid == nil then
     warn("couldn't parse nodeid under cursor")
     return

lua/pytrize/jump/util.lua

@@ -1,6 +1,9 @@
 local M = {}
 
 M.open_file = function(file)
+    -- vim.fn.execute() wraps :redir internally, suppressing the C-level
+    -- file-info message ("path" NL, NB) that :edit emits. vim.cmd.edit()
+    -- and :silent do not reliably suppress it.
     vim.fn.execute("edit " .. vim.fn.fnameescape(file))
 end
 

lua/pytrize/marks.lua

@@ -2,49 +2,27 @@ local M = {}
 
 local settings = require("pytrize.settings").settings
 
-local next_ext_ids = {}
-local ext_ids = {}
-
-local function get_ns_id()
-    return vim.api.nvim_create_namespace("pytrize")
-end
+local NS_ID = vim.api.nvim_create_namespace("pytrize")
 
+--- Clear all pytrize extmarks from a buffer.
+---@param bufnr integer Buffer number (0 for current)
 M.clear = function(bufnr)
     if bufnr == 0 then
-        bufnr = vim.fn.bufnr()
-    end
-    for _, ext_id in ipairs(ext_ids[bufnr] or {}) do
-        vim.api.nvim_buf_del_extmark(bufnr, get_ns_id(), ext_id)
-    end
-    ext_ids[bufnr] = nil
-    next_ext_ids[bufnr] = nil
-end
-
-local get_ext_id = function(bufnr)
-    if ext_ids[bufnr] == nil then
-        ext_ids[bufnr] = {}
-    end
-    if next_ext_ids[bufnr] == nil then
-        next_ext_ids[bufnr] = 1
+        bufnr = vim.api.nvim_get_current_buf()
     end
-    local next_ext_id = next_ext_ids[bufnr]
-    table.insert(ext_ids[bufnr], next_ext_id)
-    next_ext_ids[bufnr] = next_ext_id + 1
-    return next_ext_id
+    vim.api.nvim_buf_clear_namespace(bufnr, NS_ID, 0, -1)
 end
 
+--- Set a virtual text extmark.
+---@param opts { bufnr: integer, row: integer, text: string }
 M.set = function(opts)
     local bufnr = opts.bufnr
     if bufnr == 0 then
-        bufnr = vim.fn.bufnr()
+        bufnr = vim.api.nvim_get_current_buf()
     end
-    vim.api.nvim_buf_set_extmark(
-        bufnr,
-        get_ns_id(),
-        opts.row,
-        0,
-        { id = get_ext_id(bufnr), virt_text = { { opts.text, settings.highlight } } }
-    )
+    vim.api.nvim_buf_set_extmark(bufnr, NS_ID, opts.row, 0, {
+        virt_text = { { opts.text, settings.highlight } },
+    })
 end
 
 return M

lua/pytrize/nodeids.lua

@@ -6,7 +6,25 @@ local get_nodeids_path = require("pytrize.paths").get_nodeids_path
 
 local function get_raw_nodeids(rootdir)
     local nodeids_path = get_nodeids_path(rootdir)
-    return vim.fn.json_decode(vim.fn.readfile(nodeids_path))
+
+    if vim.fn.filereadable(nodeids_path) ~= 1 then
+        warn(string.format("Nodeids file not found: %s\nHave you run pytest?", nodeids_path))
+        return {}
+    end
+
+    local ok_read, content = pcall(vim.fn.readfile, nodeids_path)
+    if not ok_read then
+        warn(string.format("Failed to read nodeids file: %s", nodeids_path))
+        return {}
+    end
+
+    local ok_json, result = pcall(vim.fn.json_decode, content)
+    if not ok_json then
+        warn(string.format("Failed to parse nodeids JSON: %s", nodeids_path))
+        return {}
+    end
+
+    return result
 end
 
 M.parse_raw = function(raw_nodeid)