Merge pull request #3 from AckslD/dev

Dev

Author: AckslD

LICENSE.md

@@ -0,0 +1 @@
+Copyright (c) Axel Dahlberg. Distributed under the same terms as Vim itself. See `:help license`.

lua/pytrize/api.lua

@@ -1,51 +1,56 @@
 local M = {}
 
 M.clear = function(bufnr)
-    local marks = require('pytrize.marks')
-    marks.clear(bufnr)
+  local marks = require('pytrize.marks')
+
+  if bufnr == nil then bufnr = 0 end
+  marks.clear(bufnr)
 end
 
 M.set = function(bufnr)
-    local cs = require('pytrize.call_spec')
-    local params = require('pytrize.params')
-    local marks = require('pytrize.marks')
+  local cs = require('pytrize.call_spec')
+  local marks = require('pytrize.marks')
 
-    if bufnr == nil then
-        bufnr = 0
-    end
-    marks.clear(bufnr)
-    local param_order, call_specs = cs.get()
-    if param_order == nil then
-        return
-    end
-    local param_values = params.get_values(param_order, bufnr)
-    if param_values == nil then
-        return
-    end
-    local ext_id = 1 -- TODO why doesn't automatic assignment work
+  if bufnr == nil then bufnr = 0 end
+  marks.clear(bufnr)
+  local call_specs_per_func = cs.get_calls(bufnr)
+  if call_specs_per_func == nil then
+    return
+  end
+  for _, call_specs in pairs(call_specs_per_func) do
     for _, call_spec in ipairs(call_specs) do
-        for i, list_entry_node in ipairs(cs.list_entries(call_spec.call_node)) do
-            local param_id = params.get_id(param_values[call_spec.func_name], call_spec.params, i)
-            -- local ext_id = marks.set{ TODO
+      for _, entry_spec in ipairs(call_spec.entries) do
+        local entry_row = entry_spec.node:start()
+        marks.set{
+            bufnr = bufnr,
+            text = entry_spec.id,
+            row = entry_row,
+        }
+        for _, item_spec in ipairs(entry_spec.items) do
+          local item_row = item_spec.node:start()
+          if item_row ~= entry_row then
             marks.set{
                 bufnr = bufnr,
-                text = param_id,
-                row = list_entry_node:start(),
-                ext_id = ext_id,
+                text = item_spec.id,
+                row = item_spec.node:start(),
             }
-            ext_id = ext_id + 1 -- TODO
+          end
         end
+      end
     end
+  end
 end
 
 M.jump = function()
-    local jump = require('pytrize.jump')
-    jump.to_param_declaration()
+  local jump = require('pytrize.jump')
+
+  jump.to_param_declaration()
 end
 
 M.jump_fixture = function()
-    local jump = require('pytrize.jump')
-    jump.to_fixture_declaration()
+  local jump = require('pytrize.jump')
+
+  jump.to_fixture_declaration()
 end
 
 return M

lua/pytrize/call_spec.lua

@@ -1,99 +1,183 @@
 local M = {}
 
 local ts = vim.treesitter
-local ts_utils = require('nvim-treesitter.ts_utils')
 local ts_query = ts.query
-local reverse = require('pytrize.tables').reverse
+
 local warn = require('pytrize.warn').warn
+local tbls = require('pytrize.tables')
 
-local function get_root()
-    local parser = ts.get_parser()
+local get_root = function(bufnr)
+    local parser = ts.get_parser(bufnr)
     local tstree = parser:parse()[1]
     return tstree:root()
 end
 
-local function get_param_call_nodes()
-    local tsroot = get_root()
+local get_param_call_nodes = function(bufnr)
+    local tsroot = get_root(bufnr)
     local query = ts_query.parse_query(
         'python',
-        '(call function: (attribute) @param) (#eq? @param "pytest.mark.parametrize")'
+        -- TODO not sure why eg (#eq? @param "pytest.mark.parametrize") does not work
+        [[
+          (decorated_definition (
+            decorator (
+              call
+                function: ((attribute) @param)
+            )
+          ))
+        ]]
     )
     local nodes = {}
     for _, node, _ in query:iter_captures(tsroot) do
+      if ts.get_node_text(node, bufnr) == 'pytest.mark.parametrize' then
         table.insert(nodes, node:parent())
+      end
     end
     return nodes
 end
 
-M.get = function()
-    local calls = get_param_call_nodes()
-    local call_specs = {}
-    local param_order = {}
-    for _, call in ipairs(calls) do
-        -- Move to separate func, better way?
-        local decorated_definition = call:parent():parent()
-        if decorated_definition:type() ~= 'decorated_definition' then
-            warn("couldn't parse params")
-            return
-        end
-        local func = decorated_definition:field('definition')[1]
-        local func_name = ts_utils.get_node_text(func:field('name')[1])[1]
+local get_second_arg_node = function(call_node)
+    local arguments = call_node:field('arguments')[1]
+    return arguments:child(3)
+end
 
-        local arguments = call:field('arguments')[1]
-        local params_node = arguments:child(1)
-        if params_node:type() ~= 'string' then
-            warn("couldn't parse params")
-            return
-        end
-        local params_str = ts_utils.get_node_text(params_node)[1] -- TODO multiline strings?
-        params_str = params_str:sub(2, -2)
-        local params = vim.fn.split(params_str, [[,\s*]])  -- TODO avoid vim script?
-        table.insert(call_specs, {
-            call_node = call,
-            params = params,
-            func_name = func_name,
-        })
-        for _, param in ipairs(reverse(params)) do
-            table.insert(param_order, param)
-        end
+local get_named_children = function(node)
+  local children = {}
+  for child in node:iter_children() do
+    if child:named() and child:type() ~= 'comment' then
+      table.insert(children, child)
     end
-    return reverse(param_order), call_specs
+  end
+  return children
 end
 
-M.has_param = function(call_spec, param)
-    for _, p in ipairs(call_spec.params) do
-        if p == param then
-            return true
-        end
+local list_entries = function(call_node)
+    local list = get_second_arg_node(call_node)
+    if list:type() ~= 'list' then
+        return {}
     end
-    return false
+    return get_named_children(list)
 end
 
-local non_entry_types = {
-    ['comment'] = true,
-    -- TODO why are these even entry types?
-    ['['] = true,
-    [','] = true,
-    [']'] = true,
+local LITERALS = {
+  integer = true,
+  float = true,
+  ['true'] = true,
+  ['false'] = true,
 }
 
-M.get_second_arg_node = function(call_node)
-    local arguments = call_node:field('arguments')[1]
-    return arguments:child(3)
+local is_simple_literal = function(node)
+  return LITERALS[node:type()] ~= nil
 end
 
-M.list_entries = function(call_node)
-    local entry_nodes = {}
-    local list = M.get_second_arg_node(call_node)
-    if list:type() ~= 'list' then
-        return entry_nodes
+local get_item_id = function(entry_idx, item_node, param, bufnr)
+  if item_node:type() == 'string' then
+    local str = ts.get_node_text(item_node, bufnr)
+    local quote = str:sub(1, 1)
+    str = vim.fn.trim(str, quote):gsub('\n', '\\n')
+    return str
+  elseif is_simple_literal(item_node) then
+    return ts.get_node_text(item_node, bufnr)
+  end
+  return string.format('%s%d', param, entry_idx)
+end
+
+local get_entry = function(entry_idx, entry_node, params, bufnr)
+  if entry_node:type() ~= 'tuple' then
+    if #params == 1 then
+      return {{
+        id = get_item_id(entry_idx, entry_node, params[1], bufnr),
+        node = entry_node,
+        param = params[1],
+        idx = entry_idx,
+      }}
+    else
+      return {{
+        id = string.format('unknown (%d)', entry_idx),
+        node = entry_node,
+      }}
+    end
+  end
+  local items = {}
+  local item_nodes = get_named_children(entry_node)
+  if #params ~= #item_nodes then
+    -- TODO warn here?
+    -- warn(string.format(
+    --   'number of items in entry tuple differ from number of params, %d items and %d params (line %d in %s)',
+    --   #item_nodes,
+    --   #params,
+    --   entry_node:start() + 1,
+    --   vim.fn.bufname(bufnr)
+    -- ))
+    return nil
+  end
+  for i, param in ipairs(params) do
+    table.insert(items, {
+      id = get_item_id(entry_idx, item_nodes[i], param, bufnr),
+      node = item_nodes[i],
+      param = param,
+      idx = entry_idx,
+    })
+  end
+  return items
+end
+
+local get_entries = function(call_node, params, bufnr)
+  local entries = {}
+  for entry_idx, entry_node in ipairs(list_entries(call_node)) do
+    local items = get_entry(entry_idx - 1, entry_node, params, bufnr)
+    if items ~= nil then
+      table.insert(entries, {
+        id = table.concat(tbls.list_map(function(item) return item.id end, items), '-'),
+        items = items,
+        node = entry_node,
+      })
     end
-    for child in list:iter_children() do
-        if non_entry_types[child:type()] == nil then
-            table.insert(entry_nodes, child)
-        end
+  end
+  return entries
+end
+
+M.get_calls = function(bufnr)
+  local calls = get_param_call_nodes(bufnr or 0)
+  local call_specs = {}
+  for _, call in ipairs(calls) do
+    -- Move to separate func, better way?
+    local decorated_definition = call:parent():parent()
+    if decorated_definition:type() ~= 'decorated_definition' then
+      local row = call:start()
+      warn(string.format(
+        "couldn't parse params (line %d)\n  expected `decorated_definition`\n  got `%s`",
+        row,
+        decorated_definition:type()
+      ))
+      return
+    end
+    local func = decorated_definition:field('definition')[1]
+    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 row = call:start()
+      warn(string.format(
+        "couldn't parse params (line %d)\n  expected `string`\n  got `%s`",
+        row,
+        params_node:type()
+      ))
+      return
     end
-    return entry_nodes
+    local params_str = ts.get_node_text(params_node, bufnr)
+    params_str = params_str:sub(2, -2)
+    local params = vim.fn.split(params_str, [[,\s*]])  -- TODO avoid vim script?
+    local entries = get_entries(call, params, bufnr)
+    if call_specs[func_name] == nil then call_specs[func_name] = {} end
+    table.insert(call_specs[func_name], {
+      node = call,
+      entries = entries,
+      params = params,
+      func_name = func_name,
+    })
+  end
+  return call_specs
 end
 
 return M

lua/pytrize/input/builtin.lua

@@ -1,20 +1,20 @@
-local M = {}
-local handler = {}
-
-handler.prompt_files = function(prompt, files, callback)
-    local textlist = {prompt}
-    for i, file in ipairs(files) do
-        table.insert(textlist, string.format('%d. %s', i, file))
-    end
-    local choice = vim.fn.inputlist(textlist)
-    local file = files[choice]
-    if file then
-        callback(file)
-    end
-end
-
-M.load = function()
-    return handler
-end
-
-return M
+-- local M = {}
+-- local handler = {}
+--
+-- handler.prompt_files = function(prompt, files, callback)
+--     local textlist = {prompt}
+--     for i, file in ipairs(files) do
+--         table.insert(textlist, string.format('%d. %s', i, file))
+--     end
+--     local choice = vim.fn.inputlist(textlist)
+--     local file = files[choice]
+--     if file then
+--         callback(file)
+--     end
+-- end
+--
+-- M.load = function()
+--     return handler
+-- end
+--
+-- return M