Function argument splitting for python

Author: AndrewRadev

autoload/sj/python.vim

@@ -87,10 +87,22 @@ function! sj#python#SplitTuple()
   return s:SplitList('(.\{-})', '(', ')')
 endfunction
 
+function! sj#python#SplitArgs()
+  if search('\%#[[:keyword:].]\+(', 'e', line('.'))
+    return sj#python#SplitTuple()
+  endif
+endfunction
+
 function! sj#python#JoinTuple()
   return s:JoinList('([^()]*\s*$', '(', ')')
 endfunction
 
+function! sj#python#JoinArgs()
+  if search('\%#[[:keyword:].]\+(', 'e', line('.'))
+    return sj#python#JoinTuple()
+  endif
+endfunction
+
 function! sj#python#SplitImport()
   let import_pattern = '^\s*from \%(.*\) import \zs.*$'
 

ftplugin/python/splitjoin.vim

@@ -3,6 +3,7 @@ if !exists('b:splitjoin_split_callbacks')
         \ 'sj#python#SplitString',
         \ 'sj#python#SplitListComprehension',
         \ 'sj#python#SplitTuple',
+        \ 'sj#python#SplitArgs',
         \ 'sj#python#SplitAssignment',
         \ 'sj#python#SplitTernaryAssignment',
         \ 'sj#python#SplitDict',
@@ -18,6 +19,7 @@ if !exists('b:splitjoin_join_callbacks')
         \ 'sj#python#JoinImportWithRoundBrackets',
         \ 'sj#python#JoinMultilineString',
         \ 'sj#python#JoinTuple',
+        \ 'sj#python#JoinArgs',
         \ 'sj#python#JoinDict',
         \ 'sj#python#JoinArray',
         \ 'sj#python#JoinTernaryAssignment',

spec/plugin/python_spec.rb

@@ -64,6 +64,23 @@
     assert_file_contents 'spam = [1, [2, 3], 4]'
   end
 
+  specify "function arguments" do
+    set_file_contents 'spam = callable(one, [2, 3], four)'
+
+    vim.search 'callable'
+    split
+
+    assert_file_contents <<~EOF
+      spam = callable(one,
+                      [2, 3],
+                      four)
+    EOF
+
+    join
+
+    assert_file_contents 'spam = callable(one, [2, 3], four)'
+  end
+
   specify "imports" do
     set_file_contents <<~EOF
       def surrounding_function():