Merge pull request #77 from Konfekt/range

fix selection handling in vim_ai functions

1 files changed, 43 insertions, 47 deletions

diff --git a/autoload/vim_ai.vim b/autoload/vim_ai.vim
index dfd9c3b..ffc5fa3 100644
--- a/autoload/vim_ai.vim
+++ b/autoload/vim_ai.vim
@@ -6,6 +6,8 @@ let s:chat_py = s:plugin_root . "/py/chat.py"
 
 " remembers last command parameters to be used in AIRedoRun
 let s:last_is_selection = 0
+let s:last_firstline = 1
+let s:last_lastline = 1
 let s:last_instruction = ""
 let s:last_command = ""
 let s:last_config = {}
@@ -87,19 +89,19 @@ function! s:set_nopaste(config)
   endif
 endfunction
 
-function! s:GetCurrentLineOrSelection(is_selection)
+function! s:GetSelectionOrRange(is_selection, ...)
   if a:is_selection
     return s:GetVisualSelection()
   else
-    return trim(join(getline(a:firstline, a:lastline), "\n"))
+    return trim(join(getline(a:1, a:2), "\n"))
   endif
 endfunction
 
-function! s:SelectCurrentLineOrSelection(is_selection)
+function! s:SelectSelectionOrRange(is_selection, ...)
   if a:is_selection
     execute "normal! gv"
   else
-    execute 'normal! V'
+    execute 'normal!' . a:1 . 'GV' . a:2 . 'G'
   endif
 endfunction
 
@@ -117,22 +119,23 @@ function! s:GetVisualSelection()
 endfunction
 
 " Complete prompt
-" - is_selection - <range> parameter
 " - config       - function scoped vim_ai_complete config
 " - a:1          - optional instruction prompt
-function! vim_ai#AIRun(is_selection, config, ...) range
+" - a:2          - optional selection pending (to override g:vim_ai_is_selection_pending)
+function! vim_ai#AIRun(config, ...) range
   let l:config = vim_ai_config#ExtendDeep(g:vim_ai_complete, a:config)
-
-  let l:instruction = a:0 ? a:1 : ""
-  let l:selection = s:GetCurrentLineOrSelection(a:is_selection)
-  let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
+  let l:instruction = a:0 > 0 ? a:1 : ""
   " used for getting in Python script
-  let l:is_selection = a:is_selection
+  let l:is_selection = a:0 > 1 ? a:2 : g:vim_ai_is_selection_pending
+  let l:selection = s:GetSelectionOrRange(l:is_selection, a:firstline, a:lastline)
+  let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
 
   let s:last_command = "complete"
   let s:last_config = a:config
   let s:last_instruction = l:instruction
-  let s:last_is_selection = a:is_selection
+  let s:last_is_selection = l:is_selection
+  let s:last_firstline = a:firstline
+  let s:last_lastline = a:lastline
 
   let l:cursor_on_empty_line = empty(getline('.'))
   call s:set_paste(l:config)
@@ -147,41 +150,46 @@ function! vim_ai#AIRun(is_selection, config, ...) range
 endfunction
 
 " Edit prompt
-" - is_selection - <range> parameter
 " - config       - function scoped vim_ai_edit config
 " - a:1          - optional instruction prompt
-function! vim_ai#AIEditRun(is_selection, config, ...) range
+" - a:2          - optional selection pending (to override g:vim_ai_is_selection_pending)
+function! vim_ai#AIEditRun(config, ...) range
   let l:config = vim_ai_config#ExtendDeep(g:vim_ai_edit, a:config)
-
-  let l:instruction = a:0 ? a:1 : ""
-  let l:selection = s:GetCurrentLineOrSelection(a:is_selection)
+  let l:instruction = a:0 > 0 ? a:1 : ""
   " used for getting in Python script
-  let l:is_selection = a:is_selection
+  let l:is_selection = a:0 > 1 ? a:2 : g:vim_ai_is_selection_pending
+  let l:selection = s:GetSelectionOrRange(l:is_selection, a:firstline, a:lastline)
   let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
 
   let s:last_command = "edit"
   let s:last_config = a:config
   let s:last_instruction = l:instruction
-  let s:last_is_selection = a:is_selection
+  let s:last_is_selection = l:is_selection
+  let s:last_firstline = a:firstline
+  let s:last_lastline = a:lastline
 
   call s:set_paste(l:config)
-  call s:SelectCurrentLineOrSelection(a:is_selection)
+  call s:SelectSelectionOrRange(l:is_selection, a:firstline, a:lastline)
   execute "normal! c"
   execute "py3file " . s:complete_py
   call s:set_nopaste(l:config)
 endfunction
 
 " Start and answer the chat
-" - is_selection - <range> parameter
+" - uses_range   - true if range passed
 " - config       - function scoped vim_ai_chat config
 " - a:1          - optional instruction prompt
-function! vim_ai#AIChatRun(is_selection, config, ...) range
+function! vim_ai#AIChatRun(uses_range, config, ...) range
   let l:config = vim_ai_config#ExtendDeep(g:vim_ai_chat, a:config)
-
   let l:instruction = ""
-  let l:selection = s:GetVisualSelection()
-  " used for getting in Python script
-  let l:is_selection = a:is_selection
+  " l:is_selection used in Python script
+  if a:uses_range
+    let l:is_selection = g:vim_ai_is_selection_pending
+    let l:selection = s:GetSelectionOrRange(l:is_selection, a:firstline, a:lastline)
+  else
+    let l:is_selection = 0
+    let l:selection = ''
+  endif
   call s:set_paste(l:config)
   if &filetype != 'aichat'
     let l:chat_win_id = bufwinid(s:scratch_buffer_name)
@@ -197,15 +205,13 @@ function! vim_ai#AIChatRun(is_selection, config, ...) range
   endif
 
   let l:prompt = ""
-  if a:0 || a:is_selection
-    let l:instruction = a:0 ? a:1 : ""
+  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_instruction = l:instruction
-  let s:last_is_selection = a:is_selection
 
   execute "py3file " . s:chat_py
   call s:set_nopaste(l:config)
@@ -214,29 +220,19 @@ endfunction
 " Start a new chat
 " a:1 - optional preset shorcut (below, right, tab)
 function! vim_ai#AINewChatRun(...)
-  let l:open_conf = a:0 ? "preset_" . a:1 : g:vim_ai_chat['ui']['open_chat_command']
+  let l:open_conf = a:0 > 0 ? "preset_" . a:1 : g:vim_ai_chat['ui']['open_chat_command']
   call s:OpenChatWindow(l:open_conf)
   call vim_ai#AIChatRun(0, {})
 endfunction
 
 " Repeat last AI command
 function! vim_ai#AIRedoRun()
-  execute "normal! u"
-  if s:last_command == "complete"
-    if s:last_is_selection
-      '<,'>call vim_ai#AIRun(s:last_is_selection, s:last_config, s:last_instruction)
-    else
-      call vim_ai#AIRun(s:last_is_selection, s:last_config, s:last_instruction)
-    endif
-  endif
-  if s:last_command == "edit"
-    if s:last_is_selection
-      '<,'>call vim_ai#AIEditRun(s:last_is_selection, s:last_config, s:last_instruction)
-    else
-      call vim_ai#AIEditRun(s:last_is_selection, s:last_config, s:last_instruction)
-    endif
-  endif
-  if s:last_command == "chat"
+  undo
+  if s:last_command ==# "complete"
+    exe s:last_firstline.",".s:last_lastline . "call vim_ai#AIRun(s:last_config, s:last_instruction, s:last_is_selection)"
+  elseif s:last_command ==# "edit"
+    exe s:last_firstline.",".s:last_lastline . "call vim_ai#AIEditRun(s:last_config, s:last_instruction, s:last_is_selection)"
+  elseif s:last_command ==# "chat"
     " chat does not need prompt, all information are in the buffer already
     call vim_ai#AIChatRun(0, s:last_config)
   endif