better support for function display in strings

after/syntax/python.vim

@@ -17,6 +17,19 @@ if g:jedi#show_function_definition == 1 && has('conceal')
   hi jediFat term=bold,underline cterm=bold,underline gui=bold,underline ctermbg=0 guibg=Grey
   hi jediFunction term=NONE cterm=NONE ctermfg=6 guifg=Cyan gui=NONE ctermbg=0 guibg=Grey
 
-  " override defaults
+  " override defaults (add jediFunction to contains)
   syn match pythonComment "#.*$" contains=pythonTodo,@Spell,jediFunction
+  syn region pythonString
+      \ start=+[uU]\=\z(['"]\)+ end="\z1" skip="\\\\\|\\\z1"
+      \ contains=pythonEscape,@Spell,jediFunction
+  syn region pythonString
+      \ start=+[uU]\=\z('''\|"""\)+ end="\z1" keepend
+      \ contains=pythonEscape,pythonSpaceError,pythonDoctest,@Spell,jediFunction
+  syn region pythonRawString
+      \ start=+[uU]\=[rR]\z(['"]\)+ end="\z1" skip="\\\\\|\\\z1"
+      \ contains=@Spell,jediFunction
+  syn region pythonRawString
+      \ start=+[uU]\=[rR]\z('''\|"""\)+ end="\z1" keepend
+      \ contains=pythonSpaceError,pythonDoctest,@Spell,jediFunction
+
 endif

autoload/jedi.vim

@@ -65,8 +65,8 @@ python << PYTHONEOF
 if 1:
     try:
         show_func_def(get_script().get_in_function_call())
-    except Exception as e:
-        print e
+    except Exception:
+        print(traceback.format_exc())
 PYTHONEOF
     return ''
 endfunction
@@ -77,12 +77,16 @@ python << PYTHONEOF
 if 1:
     cursor = vim.current.window.cursor
     e = vim.eval('g:jedi#function_definition_escape')
-    regex = r'%sjedi=\([^%s]*\)%s.*%sjedi%s'.replace('%s', e)
-    vim.command(r'try | %%s/%s/\1/g | catch | endtry' % regex)
+    regex = r'%sjedi=([0-9]+), ([^%s]*)%s.*%sjedi%s'.replace('%s', e)
+    for i, line in enumerate(vim.current.buffer):
+        match = re.search(r'%s' % regex, line)
+        if match is not None:
+            vim_regex = r'\v' + regex.replace('=', r'\=') + '.{%s}' % int(match.group(1))
+            vim.command(r'try | %s,%ss/%s/\2/g | catch | endtry' % (i + 1, i + 1, vim_regex))
+            vim.command('call histdel("search", -1)')
+            vim.command('let @/ = histget("search", -1)')
     vim.current.window.cursor = cursor
 PYTHONEOF
-    call histdel("search", -1)
-    let @/ = histget("search", -1)
 endfunction
 
 

plugin/jedi.vim

@@ -168,7 +168,8 @@ def show_func_def(call_def, completion_lines=0):
     if column < 2 or row == 0:
         return  # edge cases, just ignore
 
-    row_to_replace = row - 1  # TODO check if completion menu is above or below
+    # TODO check if completion menu is above or below
+    row_to_replace = row - 1
     line = vim.eval("getline(%s)" % row_to_replace)
 
     insert_column = column - 2 # because it has stuff at the beginning
@@ -185,9 +186,14 @@ def show_func_def(call_def, completion_lines=0):
     # replace line before with cursor
     e = vim.eval('g:jedi#function_definition_escape')
     regex = "xjedi=%sx%sxjedix".replace('x', e)
-    repl = ("%s" + regex + "%s") % (line[:insert_column],
-                    line[insert_column:end_column], text, line[end_column:])
-    vim.eval('setline(%s, "%s")' % (row_to_replace, repl))
+
+    # check the replace stuff for strings, to append them (don't want to break the syntax)
+    replace = line[insert_column:end_column]
+    add = ''.join(re.findall('\\\\*["\']+', replace))  # add are all the strings
+    tup = '%s, %s' % (len(add), replace)
+    repl = ("%s" + regex + "%s") % (line[:insert_column], tup, text, add + line[end_column:])
+
+    vim.eval('setline(%s, %s)' % (row_to_replace, repr(PythonToVimStr(repl))))
 PYTHONEOF
 
 " vim: set et ts=4: