Support multiline strings in python

Author: AndrewRadev

autoload/sj/python.vim

@@ -317,6 +317,104 @@ function! sj#python#JoinTernaryAssignment()
   return 1
 endfunction
 
+function! sj#python#SplitString()
+  let char = getline('.')[col('.') - 1]
+  if char != '"' && char != "'"
+    return 0
+  endif
+
+  let string_pattern       = '\(\%(^\|[^\\]\)\zs\([''"]\)\).\{-}[^\\]\+\2'
+  let empty_string_pattern = '\%(''''\|""\)'
+
+  let lineno = line('.')
+
+  let [match_start, match_end] = sj#SearchColsUnderCursor(string_pattern)
+  if match_start <= 0
+    let [match_start, match_end] = sj#SearchColsUnderCursor(empty_string_pattern)
+    if match_start <= 0
+      return 0
+    endif
+  endif
+
+  let string    = sj#GetCols(match_start, match_end - 1)
+  let delimiter = string[0]
+  let body      = string[1:-2]
+  let indent    = indent(lineno)
+
+  if body =~ '^\(''\|""\)'
+    " then we're trying to split a string that's already multiline, ignore:
+    return 0
+  endif
+
+  if delimiter == '"'
+    if len(body) == 0
+      call sj#ReplaceMotion('vi"', '"""'.."\n"..'"""')
+    else
+      let body = substitute(body, '\\"', '"', 'g')
+      call sj#ReplaceMotion('vi"', '""'.."\n"..body.."\n".'""')
+    endif
+  elseif delimiter == "'"
+    if len(body) == 0
+      call sj#ReplaceMotion("vi'", "'''\n'''")
+    else
+      let body = substitute(body, "\\''", "'", 'g')
+      call sj#ReplaceMotion("vi'", "''\n"..body.."\n''")
+    endif
+  else
+    return 0
+  endif
+
+  if len(body) == 0
+    call sj#SetIndent(lineno + 1, lineno + 1, indent)
+  else
+    call sj#SetIndent(lineno + 1, lineno + 1, indent + shiftwidth())
+    call sj#SetIndent(lineno + 2, lineno + 2, indent)
+  endif
+
+  return 1
+endfunction
+
+function! sj#python#JoinMultilineString()
+  if sj#SearchUnderCursor('\("""\|''''''\)\s*$') <= 0
+    return 0
+  endif
+
+  let start_lineno = line('.')
+  let prefix       = getline('.')[0            : col('.') - 2]
+  let delimiter    = getline('.')[col('.') - 1 : col('.') + 2]
+
+  if search('^\s*'.delimiter, 'W') <= 0
+    return 0
+  endif
+
+  let end_lineno = line('.')
+  let suffix     = matchstr(getline(end_lineno), '^\s*'.delimiter.'\zs.*\ze')
+
+  if end_lineno - start_lineno > 1
+    let lines = sj#GetLines(start_lineno + 1, end_lineno - 1)
+    let lines = sj#TrimList(lines)
+    let body  = join(lines, " ")
+  else
+    let body = ''
+  endif
+
+  if delimiter == '"""'
+    let quote = '"'
+    let body = escape(body, '"')
+  elseif delimiter == "'''"
+    let quote = "'"
+    let body = escape(body, "'")
+  else
+    return 0
+  endif
+
+  let replacement = prefix..quote..body..quote..suffix
+  call sj#ReplaceLines(start_lineno, end_lineno, replacement)
+
+  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

@@ -963,6 +963,19 @@ Note that splitting `a, b = b, a` would not result in an expression that works
 the same way, due to the special handling by python of this case to swap two
 values.
 
+Strings ~
+>
+    db.execute('UPDATE table SET value = ?', value)
+
+    db.execute('''
+        UPDATE table SET value = ?
+    ''', value)
+<
+To avoid string-splitting overriding argument-splitting, it only triggers when
+the cursor is on the string delimiter. It's recommended to use f'/f"/F'/F" to
+position the cursor precisely before splitting.
+
+
 ==============================================================================
 R                                                                  *splitjoin-r*
 

ftplugin/python/splitjoin.vim

@@ -1,5 +1,6 @@
 if !exists('b:splitjoin_split_callbacks')
   let b:splitjoin_split_callbacks = [
+        \ 'sj#python#SplitString',
         \ 'sj#python#SplitListComprehension',
         \ 'sj#python#SplitTuple',
         \ 'sj#python#SplitAssignment',
@@ -13,6 +14,7 @@ endif
 
 if !exists('b:splitjoin_join_callbacks')
   let b:splitjoin_join_callbacks = [
+        \ 'sj#python#JoinMultilineString',
         \ 'sj#python#JoinTuple',
         \ 'sj#python#JoinDict',
         \ 'sj#python#JoinArray',

spec/plugin/python_spec.rb

@@ -271,4 +271,74 @@ def example():
       result = [x * y for x in range(1, 10) for y in range(10, 20) if x != y]
     EOF
   end
+
+  describe "strings" do
+    it "joins ''' strings into single-quoted strings" do
+      set_file_contents <<~EOF
+        string = '''
+            something, "anything"
+        '''
+      EOF
+
+      vim.search "'''"
+      join
+
+      assert_file_contents <<~EOF
+        string = 'something, "anything"'
+      EOF
+    end
+
+    it "joins \"\"\" strings into double-quoted strings" do
+      set_file_contents <<~EOF
+        string = """
+            something, 'anything'
+        """
+      EOF
+
+      vim.search '"""'
+      join
+
+      assert_file_contents <<~EOF
+        string = "something, 'anything'"
+      EOF
+    end
+
+    it "splits normal strings into multiline strings" do
+      set_file_contents 'string = "\"anything\""'
+
+      vim.search '"\"'
+      split
+
+      assert_file_contents <<~EOF
+        string = """
+            "anything"
+        """
+      EOF
+    end
+
+    it "splits empty strings into empty multiline strings" do
+      set_file_contents 'string = ""'
+
+      vim.search '"'
+      split
+
+      assert_file_contents <<~EOF
+        string = """
+        """
+      EOF
+    end
+
+    it "keeps content around the string, only splits with cursor on delimiter" do
+      set_file_contents 'string = function_call(one, "two", three)'
+
+      vim.search '"'
+      split
+
+      assert_file_contents <<~EOF
+        string = function_call(one, """
+            two
+        """, three)
+      EOF
+    end
+  end
 end