Implement ternary assignment splitjoin for python

autoload/sj.vim

@@ -555,6 +555,13 @@ function! sj#SkipSyntax(syntax_groups)
   return "synIDattr(synID(line('.'),col('.'),1),'name') =~ '".skip_pattern."'"
 endfunction
 
+function! sj#IncludeSyntax(syntax_groups)
+  let syntax_groups = a:syntax_groups
+  let include_pattern  = '\%('.join(syntax_groups, '\|').'\)'
+
+  return "synIDattr(synID(line('.'),col('.'),1),'name') !~ '".include_pattern."'"
+endfunction
+
 " Checks if the current position of the cursor is within the given limits.
 "
 function! sj#CursorBetween(start, end)

autoload/sj/python.vim

@@ -230,6 +230,93 @@ function! sj#python#JoinAssignment()
   return 1
 endfunction
 
+function! sj#python#SplitTernaryAssignment()
+  if getline('.') !~ '^\s*\%(\k\|\.\)\+\s*=\s*\S'
+    return 0
+  endif
+
+  normal! 0
+  let include_syntax = sj#IncludeSyntax(['pythonConditional'])
+
+  if sj#SearchSkip('\<if\>', include_syntax, 'W', line('.')) <= 0
+    return 0
+  endif
+  let if_col = col('.')
+
+  if sj#SearchSkip('\<else\>', include_syntax, 'W', line('.')) <= 0
+    return 0
+  endif
+
+  let else_col = col('.')
+  let line     = getline('.')
+
+  let assignment_if_true = trim(strpart(line, 0, if_col - 1))
+  let if_clause          = trim(strpart(line, if_col - 1, else_col - if_col))
+  let body_if_false      = trim(strpart(line, else_col + len('else')))
+
+  let assignment_prefix   = matchstr(assignment_if_true, '\%(\k\|\.\)\+\s*=')
+  let assignment_if_false = assignment_prefix . ' ' . body_if_false
+
+  let indent      = repeat(' ', shiftwidth())
+  let base_indent = repeat(' ', indent(line('.')))
+
+  let body = join([
+        \   base_indent . if_clause . ':',
+        \   base_indent . indent . assignment_if_true,
+        \   base_indent . 'else:',
+        \   base_indent . indent . assignment_if_false,
+        \ ], "\n")
+  call sj#ReplaceMotion('V', body)
+
+  return 1
+endfunction
+
+function! sj#python#JoinTernaryAssignment()
+  let include_syntax = sj#IncludeSyntax(['pythonConditional'])
+  let start_lineno = line('.')
+  normal! 0
+
+  if sj#SearchSkip('^\s*\zsif\>', include_syntax, 'Wc', line('.')) <= 0
+    return 0
+  endif
+  let if_line = trim(getline('.'))
+  if if_line !~ ':$'
+    return 0
+  endif
+  let if_clause = strpart(if_line, 0, len(if_line) - 1)
+
+  if search('^\s*\zs\%(\k\|\.\)\+\s*=\s*\S', 'Wc', line('.') + 1) <= 0
+    return 0
+  endif
+  let assignment_if_true = trim(getline('.'))
+  let lhs_if_true = matchstr(assignment_if_true, '^\s*\zs\%(\k\|\.\)\+\s*=')
+  let body_if_true = trim(strpart(assignment_if_true, len(lhs_if_true)))
+
+  if sj#SearchSkip('^\s*\zselse:', include_syntax, 'Wc', line('.') + 2) <= 0
+    return 0
+  endif
+  let else_line = trim(getline('.'))
+  if else_line !~ ':$'
+    return 0
+  endif
+
+  if search('^\s*\zs\%(\k\|\.\)\+\s*=\s*\S', 'Wc', line('.') + 3) <= 0
+    return 0
+  endif
+  let assignment_if_false = trim(getline('.'))
+  let lhs_if_false = matchstr(assignment_if_false, '^\s*\zs\%(\k\|\.\)\+\s*=')
+  let body_if_false = trim(strpart(assignment_if_false, len(lhs_if_false)))
+
+  if lhs_if_true != lhs_if_false
+    return 0
+  endif
+
+  let body = lhs_if_true . ' ' . body_if_true . ' ' . if_clause . ' else ' . body_if_false
+  call sj#ReplaceLines(start_lineno, start_lineno + 3, body)
+
+  return 1
+endfunction
+
 function! s:SplitList(regex, opening_char, closing_char)
   let [from, to] = sj#LocateBracesAroundCursor(a:opening_char, a:closing_char, ['pythonString'])
   if from < 0 && to < 0

doc/splitjoin.txt

@@ -882,7 +882,16 @@ Imports ~
     from foo import bar,\
             baz
 <
-Assignment ~
+Ternary assignment ~
+>
+    max_x = x1 if x1 > x2 else x2
+
+    if x1 > x2:
+        max_x = x1
+    else:
+        max_x = x2
+<
+Multiple assignment ~
 >
     a, b = foo("bar"), [one, two, three]
 

ftplugin/python/splitjoin.vim

@@ -2,6 +2,7 @@ if !exists('b:splitjoin_split_callbacks')
   let b:splitjoin_split_callbacks = [
         \ 'sj#python#SplitTuple',
         \ 'sj#python#SplitAssignment',
+        \ 'sj#python#SplitTernaryAssignment',
         \ 'sj#python#SplitDict',
         \ 'sj#python#SplitArray',
         \ 'sj#python#SplitStatement',
@@ -14,6 +15,7 @@ if !exists('b:splitjoin_join_callbacks')
         \ 'sj#python#JoinTuple',
         \ 'sj#python#JoinDict',
         \ 'sj#python#JoinArray',
+        \ 'sj#python#JoinTernaryAssignment',
         \ 'sj#python#JoinStatement',
         \ 'sj#python#JoinImport',
         \ 'sj#python#JoinAssignment',

spec/plugin/python_spec.rb

@@ -91,6 +91,31 @@
     assert_file_contents 'while True: loop()'
   end
 
+  specify "ternary clauses" do
+    set_file_contents <<~EOF
+      with indent as _:
+          max_x = x1 if x1 > x2 else x2
+    EOF
+
+    vim.search('max_x')
+    split
+
+    assert_file_contents <<~EOF
+      with indent as _:
+          if x1 > x2:
+              max_x = x1
+          else:
+              max_x = x2
+    EOF
+
+    join
+
+    assert_file_contents <<~EOF
+      with indent as _:
+          max_x = x1 if x1 > x2 else x2
+    EOF
+  end
+
   specify "splitting within a string" do
     pending "Old version on CI" if ENV['CI']