Handle nested classes and nested methods

Author: mfussenegger

lua/dap-python.lua

@@ -36,12 +36,9 @@ end
 --- Built-in are test runners for unittest, pytest and django.
 --- The key is the test runner name, the value a function to generate the
 --- module name to run and its arguments. See |dap-python.TestRunner|
----@type table<string,TestRunner>
+---@type table<string, TestRunner>
 M.test_runners = {}
 
-local function prune_nil(items)
-  return vim.tbl_filter(function(x) return x end, items)
-end
 
 local is_windows = function()
   return vim.fn.has("win32") == 1
@@ -146,29 +143,51 @@ local function get_module_path()
   end
 end
 
+
+---@return string[]
+local function flatten(...)
+  local argc = select("#", ...)
+  local result = {}
+  for i = 1, argc do
+    local arg = select(i, ...)
+    if type(arg) == "table" then
+      vim.list_extend(result, arg)
+    else
+      table.insert(result, arg)
+    end
+  end
+  return result
+end
+
+
 ---@private
-function M.test_runners.unittest(classname, methodname)
-  local path = get_module_path()
-  local test_path = table.concat(prune_nil({path, classname, methodname}), '.')
+---@param classnames string[]|string
+---@param methodname string?
+function M.test_runners.unittest(classnames, methodname)
+  local test_path = table.concat(flatten(get_module_path(), classnames, methodname), '.')
   local args = {'-v', test_path}
   return 'unittest', args
 end
 
 
 ---@private
-function M.test_runners.pytest(classname, methodname)
+---@param classnames string[]|string
+---@param methodname string?
+function M.test_runners.pytest(classnames, methodname)
   local path = vim.fn.expand('%:p')
-  local test_path = table.concat(prune_nil({path, classname, methodname}), '::')
+  local test_path = table.concat(flatten({path, classnames, methodname}), '::')
   -- -s "allow output to stdout of test"
   local args = {'-s', test_path}
   return 'pytest', args
 end
 
 
 ---@private
-function M.test_runners.django(classname, methodname)
+---@param classnames string[]|string
+---@param methodname string?
+function M.test_runners.django(classnames, methodname)
   local path = get_module_path()
-  local test_path = table.concat(prune_nil({path, classname, methodname}), '.')
+  local test_path = table.concat(flatten({path, classnames, methodname}), '.')
   local args = {'test', test_path}
   return 'django', args
 end
@@ -258,50 +277,10 @@ function M.setup(adapter_python_path, opts)
 end
 
 
-local function get_nodes(query_text, predicate)
-  local end_row = api.nvim_win_get_cursor(0)[1]
-  local ft = api.nvim_buf_get_option(0, 'filetype')
-  assert(ft == 'python', 'test_method of dap-python only works for python files, not ' .. ft)
-  local query = (vim.treesitter.query.parse
-    and vim.treesitter.query.parse(ft, query_text)
-    or vim.treesitter.parse_query(ft, query_text)
-  )
-  assert(query, 'Could not parse treesitter query. Cannot find test')
-  local parser = vim.treesitter.get_parser(0)
-  local root = (parser:parse()[1]):root()
-  local nodes = {}
-  for _, node in query:iter_captures(root, 0, 0, end_row) do
-    if predicate(node) then
-      table.insert(nodes, node)
-    end
-  end
-  return nodes
-end
-
-
-local function get_function_nodes()
-  local query_text = [[
-    (function_definition
-      name: (identifier) @name) @definition.function
-  ]]
-  return get_nodes(query_text, function(node)
-    return node:type() == 'identifier'
-  end)
-end
-
-
-local function get_class_nodes()
-  local query_text = [[
-    (class_definition
-      name: (identifier) @name) @definition.class
-  ]]
-  return get_nodes(query_text, function(node)
-    return node:type() == 'identifier'
-  end)
-end
-
-
 local function get_node_text(node)
+  if vim.treesitter.get_node_text then
+    return vim.treesitter.get_node_text(node, 0)
+  end
   local row1, col1, row2, col2 = node:range()
   if row1 == row2 then
     row2 = row2 + 1
@@ -314,24 +293,90 @@ local function get_node_text(node)
 end
 
 
-local function get_parent_classname(node)
-  local parent = node:parent()
-  while parent do
-    local type = parent:type()
-    if type == 'class_definition' then
-      for child in parent:iter_children() do
-        if child:type() == 'identifier' then
-          return get_node_text(child)
-        end
+--- Reverse list inline
+---@param list any[]
+local function reverse(list)
+  local len = #list
+  for i = 1, math.floor(len * 0.5) do
+    local opposite = len - i + 1
+    list[i], list[opposite] = list[opposite], list[i]
+  end
+end
+
+
+---@param source string|integer
+---@param subject "function"|"class"
+---@param end_row integer? defaults to cursor
+---@return TSNode[]
+function M._get_nodes(source, subject, end_row)
+  end_row = end_row or api.nvim_win_get_cursor(0)[1]
+  local query_text = [[
+    (function_definition
+      name: (identifier) @function
+    )
+
+    (class_definition
+      name: (identifier) @class
+    )
+  ]]
+  local lang = "python"
+  local query = (vim.treesitter.query.parse
+    and vim.treesitter.query.parse(lang, query_text)
+    or vim.treesitter.parse_query(lang, query_text)
+  )
+  local parser = (
+    type(source) == "number"
+    and vim.treesitter.get_parser(source, lang)
+    or vim.treesitter.get_string_parser(source --[[@as string]], lang)
+  )
+  local trees = parser:parse()
+  local root = trees[1]:root()
+  local nodes = {}
+  for id, node in query:iter_captures(root, source, 0, end_row) do
+    local capture = query.captures[id]
+    if capture == subject then
+      table.insert(nodes, node)
+    end
+  end
+  if not next(nodes) then
+    return nodes
+  end
+  if subject == "function" then
+    local result = nodes[#nodes]
+    local parent = result
+    while parent ~= nil do
+      if parent:type() == "function_definition" then
+        local ident = parent:child(1)
+        assert(ident:type() == "identifier")
+        result = ident
       end
+      parent = parent:parent()
     end
-    parent = parent:parent()
+    return { result }
+  elseif subject == "class" then
+    local last = nodes[#nodes]
+    local parent = last
+    local results = {}
+    while parent ~= nil do
+      if parent:type() == "class_definition" then
+        local ident = parent:child(1)
+        assert(ident:type() == "identifier")
+        table.insert(results, ident)
+      end
+      parent = parent:parent()
+    end
+    reverse(results)
+    return results
+  else
+    error("Expected subject 'function' or 'class', not: " .. subject)
   end
 end
 
 
+---@param classnames string[]
+---@param methodname string?
 ---@param opts DebugOpts
-local function trigger_test(classname, methodname, opts)
+local function trigger_test(classnames, methodname, opts)
   local test_runner = opts.test_runner or (M.test_runner or default_runner)
   if type(test_runner) == "function" then
     test_runner = test_runner()
@@ -342,9 +387,11 @@ local function trigger_test(classname, methodname, opts)
     return
   end
   assert(type(runner) == "function", "Test runner must be a function")
-  local module, args = runner(classname, methodname)
+  -- for BWC with custom runners which expect a string instead of a list of strings
+  local classes = #classnames == 1 and classnames[1] or classnames
+  local module, args = runner(classes, methodname)
   local config = {
-    name = table.concat(prune_nil({classname, methodname}), '.'),
+    name = table.concat(flatten(classnames, methodname), '.'),
     type = 'python',
     request = 'launch',
     module = module,
@@ -355,49 +402,51 @@ local function trigger_test(classname, methodname, opts)
 end
 
 
-local function closest_above_cursor(nodes)
-  local result
-  for _, node in pairs(nodes) do
-    if not result then
-      result = node
-    else
-      local node_row1, _, _, _ = node:range()
-      local result_row1, _, _, _ = result:range()
-      if node_row1 > result_row1 then
-        result = node
-      end
-    end
-  end
-  return result
-end
-
-
 --- Run test class above cursor
 ---@param opts? DebugOpts See |dap-python.DebugOpts|
 function M.test_class(opts)
   opts = vim.tbl_extend('keep', opts or {}, default_test_opts)
-  local class_node = closest_above_cursor(get_class_nodes())
-  if not class_node then
-    print('No suitable test class found')
+  local candidates = M._get_nodes(0, "class")
+  if not candidates then
+    print('No test class found near cursor')
     return
   end
-  local class = get_node_text(class_node)
-  trigger_test(class, nil, opts)
+  local names = vim.tbl_map(get_node_text, candidates)
+  trigger_test(names, nil, opts)
+end
+
+
+---@param node TSNode
+---@result TSNode[]
+local function get_parent_classes(node)
+  local parent = node:parent()
+  local result = {}
+  while parent ~= nil do
+    if parent:type() == "class_definition" then
+      local ident = parent:child(1)
+      assert(ident and ident:type() == "identifier")
+      table.insert(result, ident)
+    end
+    parent = parent:parent()
+  end
+  reverse(result)
+  return result
 end
 
 
 --- Run the test method above cursor
 ---@param opts? DebugOpts See |dap-python.DebugOpts|
 function M.test_method(opts)
   opts = vim.tbl_extend('keep', opts or {}, default_test_opts)
-  local function_node = closest_above_cursor(get_function_nodes())
-  if not function_node then
-    print('No suitable test method found')
+  local functions = M._get_nodes(0, "function")
+  if not functions then
+    print('No test method found near cursor')
     return
   end
-  local class = get_parent_classname(function_node)
-  local function_name = get_node_text(function_node)
-  trigger_test(class, function_name, opts)
+  local fn = functions[1]
+  local parent_classes = get_parent_classes(fn)
+  local classnames = vim.tbl_map(get_node_text, parent_classes)
+  trigger_test(classnames, get_node_text(fn), opts)
 end
 
 
@@ -414,6 +463,7 @@ local function remove_indent(lines)
     end
   end
   if offset > 1 then
+    assert(offset)
     return vim.tbl_map(function(x) return string.sub(x, offset) end, lines)
   else
     return lines
@@ -479,7 +529,7 @@ end
 ---@field pythonPath string|nil Path to python interpreter. Uses interpreter from `VIRTUAL_ENV` environment variable or `adapter_python_path` by default
 
 
----@alias TestRunner fun(classname: string, methodname: string):string, string[]
+---@alias TestRunner fun(classname: string|string[], methodname: string?):string, string[]
 
 ---@alias DebugpyConsole "internalConsole"|"integratedTerminal"|"externalTerminal"|nil