handle paste mode in finally block

1 files changed, 48 insertions, 30 deletions

diff --git a/autoload/vim_ai.vim b/autoload/vim_ai.vim
index 8d5b3bc..a62ea89 100644
--- a/autoload/vim_ai.vim
+++ b/autoload/vim_ai.vim
@@ -78,15 +78,21 @@ function! s:OpenChatWindow(open_conf)
   execute l:open_cmd
 endfunction
 
+
+let s:is_handling_paste_mode = 0
+
 function! s:set_paste(config)
-  if !a:config['ui']['paste_mode'] | return | endif
-  if &paste | return | endif
-  setlocal paste
-  augroup AiPaste
-    autocmd!
-    autocmd ModeChanged i:* exe 'set nopaste'
-    autocmd! AiPaste InsertLeave
-  augroup END
+  if !&paste && a:config['ui']['paste_mode']
+    let s:is_handling_paste_mode = 1
+    setlocal paste
+  endif
+endfunction
+
+function! s:set_nopaste(config)
+  if s:is_handling_paste_mode
+    setlocal nopaste
+    let s:is_handling_paste_mode = 0
+  endif
 endfunction
 
 function! s:GetSelectionOrRange(is_selection, ...)
@@ -144,14 +150,18 @@ function! vim_ai#AIRun(config, ...) range
   let s:last_lastline = a:lastline
 
   let l:cursor_on_empty_line = empty(getline('.'))
-  call s:set_paste(l:config)
-  if l:cursor_on_empty_line
-    execute "normal! " . a:lastline . "GA"
-  else
-    execute "normal! " . a:lastline . "Go"
-  endif
-  execute "py3file " . s:complete_py
-  execute "normal! " . a:lastline . "G"
+  try
+    call s:set_paste(l:config)
+    if l:cursor_on_empty_line
+      execute "normal! " . a:lastline . "GA"
+    else
+      execute "normal! " . a:lastline . "Go"
+    endif
+    execute "py3file " . s:complete_py
+    execute "normal! " . a:lastline . "G"
+  finally
+    call s:set_nopaste(l:config)
+  endtry
 endfunction
 
 " Edit prompt
@@ -178,10 +188,14 @@ function! vim_ai#AIEditRun(config, ...) range
   let s:last_firstline = a:firstline
   let s:last_lastline = a:lastline
 
-  call s:set_paste(l:config)
-  call s:SelectSelectionOrRange(l:is_selection, a:firstline, a:lastline)
-  execute "normal! c"
-  execute "py3file " . s:complete_py
+  try
+    call s:set_paste(l:config)
+    call s:SelectSelectionOrRange(l:is_selection, a:firstline, a:lastline)
+    execute "normal! c"
+    execute "py3file " . s:complete_py
+  finally
+    call s:set_nopaste(l:config)
+  endtry
 endfunction
 
 function! s:ReuseOrCreateChatWindow(config)
@@ -234,20 +248,24 @@ function! vim_ai#AIChatRun(uses_range, config, ...) range
     let l:is_selection = 0
     let l:selection = ''
   endif
-  call s:set_paste(l:config)
+  try
+    call s:set_paste(l:config)
 
-  call s:ReuseOrCreateChatWindow(l:config)
+    call s:ReuseOrCreateChatWindow(l:config)
 
-  let l:prompt = ""
-  if a:0 > 0 || a:uses_range
-    let l:instruction = a:0 > 0 ? a:1 : ""
-    let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
-  endif
+    let l:prompt = ""
+    if a:0 > 0 || a:uses_range
+      let l:instruction = a:0 > 0 ? a:1 : ""
+      let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
+    endif
 
-  let s:last_command = "chat"
-  let s:last_config = a:config
+    let s:last_command = "chat"
+    let s:last_config = a:config
 
-  execute "py3file " . s:chat_py
+    execute "py3file " . s:chat_py
+  finally
+    call s:set_nopaste(l:config)
+  endtry
 endfunction
 
 " Start a new chat