feat(python): add "return" patterns to @indent_end

`return ...` is usually the last line of a
function/statement, so mark some possible patterns of `return` with
`@indent_end` to dedent the line after it

3 files changed, 29 insertions, 0 deletions

diff --git a/queries/python/indents.scm b/queries/python/indents.scm
index c02879712..cc13587ec 100644
--- a/queries/python/indents.scm
+++ b/queries/python/indents.scm
@@ -72,6 +72,15 @@
 (list_pattern "]" @indent_end)
 
 
+(return_statement
+  [
+    (_) @indent_end 
+    (_ (_) @indent_end .)
+    (attribute 
+      attribute: (_) @indent_end)
+    "return" @indent_end
+  ] .)
+
 [
   ")"
   "]"
diff --git a/tests/indent/python/return_dedent.py b/tests/indent/python/return_dedent.py
new file mode 100644
index 000000000..44d6219de
--- /dev/null
+++ b/tests/indent/python/return_dedent.py
@@ -0,0 +1,16 @@
+def a():
+    return
+
+def a():
+    return True
+
+def a():
+    return (1, 2, 3)
+
+def a():
+    return x.y.z
+
+def a():
+    return (
+        1, 2, 3
+    )
diff --git a/tests/indent/python_spec.lua b/tests/indent/python_spec.lua
index 6b417a37f..c400de542 100644
--- a/tests/indent/python_spec.lua
+++ b/tests/indent/python_spec.lua
@@ -74,5 +74,9 @@ describe("indent Python:", function()
     run:new_line("line_after_indent.py", { on_line = 49, text = "x", indent = 0 })
     run:new_line("line_after_indent.py", { on_line = 55, text = "x", indent = 4 })
     run:new_line("line_after_indent.py", { on_line = 63, text = "x", indent = 4 })
+
+    for _, line in ipairs { 2, 5, 8, 11, 16 } do
+      run:new_line("return_dedent.py", { on_line = line, text = "x", indent = 0 })
+    end
   end)
 end)