Improve indent for line continuation

indent/python.vim

@@ -1,7 +1,7 @@
 " Vim indent file
 " Language:    Python
 " Author:      Akinori Hattori <hattya@gmail.com>
-" Last Change: 2023-06-24
+" Last Change: 2023-07-09
 " License:     MIT License
 
 if exists('b:did_indent')
@@ -129,7 +129,15 @@ function! GetPEP8PythonIndent() abort
 
   let ind = indent(lnum)
   let ll = join(buf, ' ')
-  if ll =~# '\v:\s*%(#.*)=$'
+  if getline(v:lnum - 1) =~# s:lcont
+    " line continuation
+    let kw = s:matchkw(ll)
+    if kw ==# ''
+      let ind += s:cont()
+    else
+      let ind += len(substitute(kw, '\s\+', ' ', '')) + 1
+    endif
+  elseif ll =~# '\v:\s*%(#.*)=$'
     " compound statement
     let ind += shiftwidth()
   elseif ll =~# s:dedent
@@ -138,17 +146,6 @@ function! GetPEP8PythonIndent() abort
   elseif ll =~# s:ellipsis
     " ellipsis
     let ind -= shiftwidth()
-  elseif getline(v:lnum - 1) =~# s:lcont
-    " line continuation
-    if s:is_compound_stmt(ll)
-      if ll =~# '\v^\s*<with>'
-        let ind += len('with') + 1
-      else
-        let ind += s:ml_stmt() ? shiftwidth() * 2 : shiftwidth()
-      endif
-    else
-      let ind += s:cont()
-    endif
   elseif colon
     " : is at EOL
     return -1
@@ -182,6 +179,10 @@ function! s:is_compound_stmt(str, ...) abort
   return (a:0 && a:1 && a:str =~# '\v^\s*<%(def|class)>') || a:str =~# '\v^\s*<%(if|elif|while|for|except|with)>'
 endfunction
 
+function! s:matchkw(str) abort
+  return matchstr(a:str, '\v^\s*\zs<%(assert|del|return|yield%(\s+from)=|raise|import|from|global|nonlocal|if|elif|while|for|except|with|def|class)>')
+endfunction
+
 function! s:cont() abort
   return eval(s:getvar('python_indent_continue', shiftwidth()))
 endfunction

test/indent.vimspec

@@ -352,6 +352,91 @@ Describe filetype indent
     End
   End
 
+  Describe simple statement
+    It indents the assert statement
+      let in  = "assert True \\\<CR>or False"
+      let out = [
+      \  'assert True \',
+      \  '       or False',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+
+    It indents the del statement
+      let in  = "del spam, \\\<CR>eggs"
+      let out = [
+      \  'del spam, \',
+      \  '    eggs',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+
+    It indents the return statement
+      let in  = "return True \\\<CR>or False"
+      let out = [
+      \  'return True \',
+      \  '       or False',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+
+    It indents the yield statement
+      let in  = "yield spam \\\<CR>.eggs()\<CR>"
+      let in .= "yield from ham \\\<CR>.toast()\<CR>"
+      let in .= "yield \\\<CR>from \\\<CR>beans \\\<CR>.bacon()"
+      let out = [
+      \  'yield spam \',
+      \  '      .eggs()',
+      \  'yield from ham \',
+      \  '           .toast()',
+      \  'yield \',
+      \  '      from \',
+      \  '           beans \',
+      \  '           .bacon()',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+
+    It indents the raise statement
+      let in  = "raise Exception() \\\<CR>.with_traceback(tb)"
+      let out = [
+      \  'raise Exception() \',
+      \  '      .with_traceback(tb)',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+
+    It indents the import statement
+      let in  = "import spam \\\<CR>.eggs\<CR>"
+      let in .= "from ham \\\<CR>.toast"
+      let out = [
+      \  'import spam \',
+      \  '       .eggs',
+      \  'from ham \',
+      \  '     .toast',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+
+    It indents the global statement
+      let in  = "global spam, \\\<CR>eggs"
+      let out = [
+      \  'global spam, \',
+      \  '       eggs',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+
+    It indents the nonlocal statement
+      let in  = "nonlocal spam, \\\<CR>eggs"
+      let out = [
+      \  'nonlocal spam, \',
+      \  '         eggs',
+      \]
+      Assert Equals(Insert(in), Buffer(out))
+    End
+  End
+
   Describe compound statement
     Describe if statement
       It aligns with left brackets
@@ -421,24 +506,10 @@ Describe filetype indent
         let in .= "elif False and \\\<CR>True:"
         let out = [
         \  'if True and \',
-        \  '    False:',
-        \  '    pass',
-        \  'elif False and \',
-        \  '    True:',
-        \]
-        Assert Equals(Insert(in), Buffer(out))
-      End
-
-      It increases the indent level by 2 (line continuation)
-        let b:python_indent_multiline_statement = 1
-        let in  = "if True and \\\<CR>False:\<CR>pass\<CR>"
-        let in .= "elif False and \\\<CR>True:"
-        let out = [
-        \  'if True and \',
-        \  '        False:',
+        \  '   False:',
         \  '    pass',
         \  'elif False and \',
-        \  '        True:',
+        \  '     True:',
         \]
         Assert Equals(Insert(in), Buffer(out))
       End
@@ -526,17 +597,7 @@ Describe filetype indent
         let in  = "while True and \\\<CR>False:"
         let out = [
         \  'while True and \',
-        \  '    False:',
-        \]
-        Assert Equals(Insert(in), Buffer(out))
-      End
-
-      It increases the indent level by 2 (line continuation)
-        let b:python_indent_multiline_statement = 1
-        let in  = "while True and \\\<CR>False:"
-        let out = [
-        \  'while True and \',
-        \  '        False:',
+        \  '      False:',
         \]
         Assert Equals(Insert(in), Buffer(out))
       End
@@ -610,16 +671,6 @@ Describe filetype indent
         Assert Equals(Insert(in), Buffer(out))
       End
 
-      It increases the indent level by 2 (line continuation)
-        let b:python_indent_multiline_statement = 1
-        let in  = "for _ in True, \\\<CR>False:"
-        let out = [
-        \  'for _ in True, \',
-        \  '        False:',
-        \]
-        Assert Equals(Insert(in), Buffer(out))
-      End
-
       It indents the else clause
         let in  = "for _ in range(1):\<CR>"
         let in .= "for _ in range(2):\<CR>continue\<CR>"
@@ -699,20 +750,7 @@ Describe filetype indent
         \  'try:',
         \  '    ...',
         \  'except TypeError, \',
-        \  '    ValueError:',
-        \]
-        Assert Equals(Insert(in), Buffer(out))
-      End
-
-      It increases the indent level by 2 (line continuation)
-        let b:python_indent_multiline_statement = 1
-        let in  = "try:\<CR>...\<CR>"
-        let in .= "except TypeError, \\\<CR>ValueError:"
-        let out = [
-        \  'try:',
-        \  '    ...',
-        \  'except TypeError, \',
-        \  '        ValueError:',
+        \  '       ValueError:',
         \]
         Assert Equals(Insert(in), Buffer(out))
       End