fix: change folding algorithm to fix Python indents

4 files changed, 85 insertions, 17 deletions

diff --git a/lua/nvim-treesitter/indent.lua b/lua/nvim-treesitter/indent.lua
index c3e4c0888..cf819d05b 100644
--- a/lua/nvim-treesitter/indent.lua
+++ b/lua/nvim-treesitter/indent.lua
@@ -137,7 +137,13 @@ function M.get_indent(lnum)
 
   while node do
     -- do 'autoindent' if not marked as @indent
-    if not q.indent[node:id()] and q.auto[node:id()] and node:start() < lnum - 1 and lnum - 1 <= node:end_() then
+    if
+      not q.indent[node:id()]
+      and not q.aligned_indent[node:id()]
+      and q.auto[node:id()]
+      and node:start() < lnum - 1
+      and lnum - 1 <= node:end_()
+    then
       return -1
     end
 
@@ -171,7 +177,7 @@ function M.get_indent(lnum)
       should_process
       and (
         q.indent[node:id()]
-        and (srow ~= erow or is_in_err)
+        and (srow ~= erow or is_in_err or q.indent[node:id()].immediate_indent)
         and (srow ~= lnum - 1 or q.indent[node:id()].start_at_same_line)
       )
     then
@@ -183,12 +189,16 @@ function M.get_indent(lnum)
     if q.aligned_indent[node:id()] and srow ~= erow and (srow ~= lnum - 1) then
       local metadata = q.aligned_indent[node:id()]
       local o_delim_node, is_last_in_line ---@type TSNode|nil, boolean|nil
+      local c_delim_node ---@type TSNode|nil
       if metadata.delimiter then
         ---@type string
         local opening_delimiter = metadata.delimiter and metadata.delimiter:sub(1, 1)
         o_delim_node, is_last_in_line = find_delimiter(bufnr, node, opening_delimiter)
+        local closing_delimiter = metadata.delimiter and metadata.delimiter:sub(2, 2)
+        c_delim_node, _ = find_delimiter(bufnr, node, closing_delimiter)
       else
         o_delim_node = node
+        c_delim_node = node
       end
 
       if o_delim_node then
@@ -196,8 +206,29 @@ function M.get_indent(lnum)
           -- hanging indent (previous line ended with starting delimiter)
           indent = indent + indent_size * 1
         else
-          local _, o_scol = o_delim_node:start()
-          return math.max(indent, 0) + o_scol + (metadata.increment or 1)
+          local o_srow, o_scol = o_delim_node:start()
+          local final_line_indent = false
+          if c_delim_node then
+            local c_srow, _ = c_delim_node:start()
+            if c_srow ~= o_srow and c_srow == lnum - 1 then
+              -- delims end on current line, and are not open and closed same line.
+              -- final_line_indent controls this behavior, for example this is not desirable
+              -- for a tuple.
+              final_line_indent = metadata.final_line_indent or false
+            end
+          end
+          if final_line_indent then
+            -- last line must be indented more in cases where
+            -- it would be same indent as next line
+            local aligned_indent = o_scol + (metadata.increment or 1)
+            if aligned_indent <= indent then
+              return indent + indent_size * 1
+            else
+              return aligned_indent
+            end
+          else
+            return o_scol + (metadata.increment or 1)
+          end
         end
       end
     end
diff --git a/queries/python/indents.scm b/queries/python/indents.scm
index 3032cedfb..da37c1f30 100644
--- a/queries/python/indents.scm
+++ b/queries/python/indents.scm
@@ -3,11 +3,6 @@
   (dictionary)
   (set)
 
-  (if_statement)
-  (for_statement)
-  (while_statement)
-  (with_statement)
-  (try_statement)
   (import_from_statement)
 
   (parenthesized_expression)
@@ -21,24 +16,45 @@
   (binary_operator)
 
   (lambda)
-  (function_definition)
-  (class_definition)
 
   (concatenated_string)
 ] @indent
-  
+
+((for_statement) @indent
+ (#set! "immediate_indent" 1))
+((if_statement) @indent
+ (#set! "immediate_indent" 1))
+((while_statement) @indent
+ (#set! "immediate_indent" 1))
+((try_statement) @indent
+ (#set! "immediate_indent" 1))
+((ERROR "try" ":") @indent
+ (#set! "immediate_indent" 1))
+((function_definition) @indent
+ (#set! "immediate_indent" 1))
+((class_definition) @indent
+ (#set! "immediate_indent" 1))
+((with_statement) @indent
+ (#set! "immediate_indent" 1))
+
 (if_statement
   condition: (parenthesized_expression) @aligned_indent
   (#set! "delimiter" "()")
+  (#set! "final_line_indent" 1) ; parenthesized_expression already indented
+)
+(while_statement
+  condition: (parenthesized_expression) @aligned_indent
+  (#set! "delimiter" "()")
+  (#set! "final_line_indent" 1) ; parenthesized_expression already indented
 )
+
 ((ERROR "(" . (_)) @aligned_indent
  (#set! "delimiter" "()"))
 ((argument_list) @aligned_indent
  (#set! "delimiter" "()"))
-((argument_list) @aligned_indent
- (#set! "delimiter" "()"))
 ((parameters) @aligned_indent
- (#set! "delimiter" "()"))
+ (#set! "delimiter" "()")
+ (#set! "final_line_indent" 1))
 ((tuple) @aligned_indent
  (#set! "delimiter" "()"))
 
diff --git a/tests/indent/python/control_flow.py b/tests/indent/python/control_flow.py
index c8cd2541b..7ec02e3ff 100644
--- a/tests/indent/python/control_flow.py
+++ b/tests/indent/python/control_flow.py
@@ -20,3 +20,11 @@ except:
     pass
 finally:
     pass
+
+while (a > 4 and
+       b > 5):
+    pass
+
+try:
+    pass
+
diff --git a/tests/indent/python_spec.lua b/tests/indent/python_spec.lua
index b8856b3a0..b1e94cbf8 100644
--- a/tests/indent/python_spec.lua
+++ b/tests/indent/python_spec.lua
@@ -19,12 +19,25 @@ describe("indent Python:", function()
     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 })
     run:new_line("basic_blocks.py", { on_line = 7, text = "x += 1", indent = 4 })
-    run:new_line("basic_blocks.py", { on_line = 11, text = "x += 1", indent = 8 })
+    run:new_line("basic_blocks.py", { on_line = 10, text = "x += 1", indent = 8 })
+    run:new_line("basic_blocks.py", { on_line = 14, 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 = 2, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 4, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 6, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 9, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 12, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 15, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 18, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 20, text = "x = 4", indent = 4 })
     run:new_line("control_flow.py", { on_line = 22, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 24, text = "c < 6 and", indent = 7 })
+    run:new_line("control_flow.py", { on_line = 26, text = "x = 4", indent = 4 })
+    run:new_line("control_flow.py", { on_line = 29, text = "x = 4", indent = 4 })
+    run:new_line("branches.py", { on_line = 25, text = "x > 9 and", indent = 4 })
+    run:new_line("branches.py", { on_line = 29, text = "and x > 9", 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(