feat(indent): support @aligned_indent for python

CONTRIBUTING.md

@@ -269,11 +269,12 @@ the node describing the language and `@content` to describe the injection region
 ### Indents
 
 ```
-@indent ; Indent children when matching this node
-@dedent ; Dedent children when matching this node
-@branch ; Dedent itself when matching this node
-@ignore ; Do not indent in this node
-@auto   ; Behaves like 'autoindent' buffer option
+@indent         ; Indent children when matching this node
+@aligned_indent ; Behaves like python aligned/hanging indent
+@dedent         ; Dedent children when matching this node
+@branch         ; Dedent itself when matching this node
+@ignore         ; Do not indent in this node
+@auto           ; Behaves like 'autoindent' buffer option
 ```
 
 [Zulip]: nvim-treesitter.zulipchat.com

lua/nvim-treesitter/indent.lua

@@ -12,6 +12,19 @@ local function get_last_node_at_line(root, lnum)
   return root:descendant_for_range(lnum - 1, col, lnum - 1, col)
 end
 
+local function get_matching_prev_sibling(anchor, start, matcher)
+  local start_row, start_col = start[1], start[2]
+  local node = anchor:descendant_for_range(start_row, start_col, start_row, start_col)
+  local pos = 1
+  -- TODO: reconsider this 999 limit or do something differently in future.
+  --       if anchor has more than 999 children, this would not work.
+  while pos < 999 and node and not matcher(node) do
+    node = node:prev_sibling()
+    pos = pos + 1
+  end
+  return node, pos
+end
+
 local M = {}
 
 local get_indents = tsutils.memoize_by_buf_tick(function(bufnr, root, lang)
@@ -21,6 +34,7 @@ local get_indents = tsutils.memoize_by_buf_tick(function(bufnr, root, lang)
     dedent = {},
     branch = {},
     ignore = {},
+    aligned_indent = {},
   }
 
   for name, node, metadata in queries.iter_captures(bufnr, "indents", root, lang) do
@@ -101,6 +115,30 @@ function M.get_indent(lnum)
       is_processed = true
     end
 
+    if q.aligned_indent[node:id()] and srow ~= erow then
+      local metadata = q.aligned_indent[node:id()]
+      local opening_delimiter = metadata.delimiter:sub(1, 1)
+      local o_delim_node, pos = get_matching_prev_sibling(node, { srow, #vim.fn.getline(srow + 1) - 1 }, function(n)
+        return n:type() == opening_delimiter
+      end)
+
+      if o_delim_node then
+        if pos == 1 then
+          -- hanging indent (previous line ended with starting delimiter)
+          indent = indent + indent_size * 1
+        else
+          local _, o_scol = o_delim_node:start()
+          local aligned_indent = math.max(indent, 0) + o_scol
+          if indent > 0 then
+            indent = aligned_indent
+          else
+            indent = aligned_indent + 1 -- extra space for starting delimiter
+          end
+          is_processed = true
+        end
+      end
+    end
+
     is_processed_by_row[srow] = is_processed_by_row[srow] or is_processed
 
     node = node:parent()

queries/python/indents.scm

@@ -1,6 +1,5 @@
 [
   (list)
-  (tuple)
   (dictionary)
   (set)
 
@@ -19,8 +18,6 @@
 
   (tuple_pattern)
   (list_pattern)
-  (argument_list)
-  (parameters)
   (binary_operator)
 
   (lambda)
@@ -30,6 +27,20 @@
   (concatenated_string)
 ] @indent
 
+  
+(if_statement
+  condition: (parenthesized_expression) @aligned_indent
+  (#set! "delimiter" "()")
+)
+((argument_list) @aligned_indent
+ (#set! "delimiter" "()"))
+((argument_list) @aligned_indent
+ (#set! "delimiter" "()"))
+((parameters) @aligned_indent
+ (#set! "delimiter" "()"))
+((tuple) @aligned_indent
+ (#set! "delimiter" "()"))
+
 [
   ")"
   "]"

tests/indent/python/branches.py

@@ -23,5 +23,9 @@ foo(
     b)
 
 if (a and
-        b):
+    b):
+    pass
+
+if (a
+        and b):
     pass

tests/indent/python_spec.lua

@@ -1,5 +1,4 @@
 local Runner = require("tests.indent.common").Runner
-local XFAIL = require("tests.indent.common").XFAIL
 
 local run = Runner:new(it, "tests/indent/python", {
   tabstop = 4,
@@ -12,16 +11,13 @@ describe("indent Python:", function()
   describe("whole file:", function()
     run:whole_file(".", {
       expected_failures = {
-        "./aligned_indent.py",
         "./branches.py",
-        "./hanging_indent.py",
-        "./nested_collections.py",
       },
     })
   end)
 
   describe("new line:", function()
-    run:new_line("aligned_indent.py", { on_line = 1, text = "arg3,", indent = 19 }, "xfail", XFAIL)
+    run:new_line("aligned_indent.py", { on_line = 1, text = "arg3,", indent = 19 })
     run:new_line("basic_blocks.py", { on_line = 1, text = "wait,", indent = 4 })
     run:new_line("basic_blocks.py", { on_line = 6, text = "x += 1", indent = 4 })
     run:new_line("basic_blocks.py", { on_line = 10, text = "x += 1", indent = 8 })
@@ -29,8 +25,8 @@ describe("indent Python:", function()
     run:new_line("basic_blocks.py", { on_line = 11, text = "x += 1", indent = 8 })
     run:new_line("basic_collections.py", { on_line = 3, text = "4,", indent = 4 })
     run:new_line("comprehensions.py", { on_line = 8, text = "if x != 2", indent = 4 })
-    run:new_line("control_flow.py", { on_line = 23, text = "x = 4", indent = 4 }, "expected failure", XFAIL)
-    run:new_line("hanging_indent.py", { on_line = 1, text = "arg0,", indent = 8 }, "expected failure", XFAIL)
+    run:new_line("control_flow.py", { on_line = 22, text = "x = 4", indent = 4 })
+    run:new_line("hanging_indent.py", { on_line = 1, text = "arg0,", indent = 8 })
     run:new_line("hanging_indent.py", { on_line = 5, text = "0,", indent = 4 })
     run:new_line("join_lines.py", { on_line = 1, text = "+ 1 \\", indent = 4 })
     run:new_line("join_lines.py", { on_line = 4, text = "+ 1 \\", indent = 4 })
@@ -38,7 +34,7 @@ describe("indent Python:", function()
     run:new_line("nested_collections.py", { on_line = 5, text = "0,", indent = 12 })
     run:new_line("nested_collections.py", { on_line = 6, text = ",0", indent = 12 })
     run:new_line("nested_collections.py", { on_line = 29, text = "[1, 2],", indent = 12 })
-    run:new_line("nested_collections.py", { on_line = 39, text = "0,", indent = 5 }, "expected failure", XFAIL)
+    run:new_line("nested_collections.py", { on_line = 39, text = "0,", indent = 5 })
     run:new_line("strings.py", { on_line = 14, text = "x", indent = 4 })
     run:new_line("strings.py", { on_line = 15, text = "x", indent = 0 })
     run:new_line("strings.py", { on_line = 16, text = "x", indent = 8 })