Merge pull request #77 from Konfekt/range

fix selection handling in vim_ai functions

4 files changed, 66 insertions, 61 deletions

diff --git a/README.md b/README.md
index 5d7f26c..1bc1f34 100644
--- a/README.md
+++ b/README.md
@@ -377,7 +377,6 @@ To create a custom command, you can call `AIRun`, `AIEditRun` and `AIChatRun` fu
 function! GitCommitMessageFn()
   let l:diff = system('git --no-pager diff --staged')
   let l:prompt = "generate a short commit message from the diff below:\n" . l:diff
-  let l:range = 0
   let l:config = {
   \  "engine": "chat",
   \  "options": {
@@ -386,7 +385,7 @@ function! GitCommitMessageFn()
   \    "temperature": 1,
   \  },
   \}
-  call vim_ai#AIRun(l:range, l:config, l:prompt)
+  call vim_ai#AIRun(l:config, l:prompt)
 endfunction
 command! GitCommitMessage call GitCommitMessageFn()
 
@@ -398,9 +397,9 @@ function! CodeReviewFn(range) range
   \    "initial_prompt": ">>> system\nyou are a clean code expert",
   \  },
   \}
-  '<,'>call vim_ai#AIChatRun(a:range, l:config, l:prompt)
+  exe a:firstline.",".a:lastline . "call vim_ai#AIChatRun(a:range, l:config, l:prompt)"
 endfunction
-command! -range CodeReview <line1>,<line2>call CodeReviewFn(<range>)
+command! -range=0 CodeReview <line1>,<line2>call CodeReviewFn(<count>)
 ```
 
 ## Contributing
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
diff --git a/doc/vim-ai.txt b/doc/vim-ai.txt
index 34765d7..72535b3 100644
--- a/doc/vim-ai.txt
+++ b/doc/vim-ai.txt
@@ -188,7 +188,6 @@ To create custom commands, call `AIRun`, `AIEditRun` and `AIChatRun` functions:
   function! AIPromptCommitMessageFn()
     let l:diff = system('git diff --staged')
     let l:prompt = "generate a short commit message from the diff below:\n" . l:diff
-    let l:range = 0
     let l:config = {
     \  "engine": "chat",
     \  "options": {
@@ -197,12 +196,12 @@ To create custom commands, call `AIRun`, `AIEditRun` and `AIChatRun` functions:
     \    "temperature": 1,
     \  },
     \}
-    call vim_ai#AIRun(l:range, l:config, l:prompt)
+    call vim_ai#AIRun(l:config, l:prompt)
   endfunction
-  command! AIPromptCommitMessage call AIPromptCommitMessageFn()
+  command! AIPromptCommitMessage call AIPromptCommitMessageFn(range)
 
   " custom command that provides a code review for selected code block
-  function! AIPromptCodeReviewFn(range) range
+  function! AIPromptCodeReviewFn() range
     let l:prompt = "programming syntax is " . &filetype . ", review the code below"
     let l:config = {
     \  "options": {
@@ -211,7 +210,7 @@ To create custom commands, call `AIRun`, `AIEditRun` and `AIChatRun` functions:
     \}
     '<,'>call vim_ai#AIChatRun(a:range, l:config, l:prompt)
   endfunction
-  command! -range AIPromptCodeReview <line1>,<line2>call AIPromptCodeReviewFn(<range>)
+  command! -range=0 AIPromptCodeReview <line1>,<line2>call AIPromptCodeReviewFn(<count>)
 
 ABOUT                                           *vim-ai-about*
 
diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim
index 1d03177..5d156c3 100644
--- a/plugin/vim-ai.vim
+++ b/plugin/vim-ai.vim
@@ -4,8 +4,19 @@ if !has('python3')
   finish
 endif
 
-command! -range -nargs=? AI <line1>,<line2>call vim_ai#AIRun(<range>, {}, <f-args>)
-command! -range -nargs=? AIEdit <line1>,<line2>call vim_ai#AIEditRun(<range>, {}, <f-args>)
-command! -range -nargs=? AIChat <line1>,<line2>call vim_ai#AIChatRun(<range>, {}, <f-args>)
-command! -nargs=? AINewChat call vim_ai#AINewChatRun(<f-args>)
-command! AIRedo call vim_ai#AIRedoRun()
+" detect if a visual selection is pending: https://stackoverflow.com/a/20133772
+let g:vim_ai_is_selection_pending = 0
+augroup vim_ai
+  autocmd!
+    autocmd CursorMoved *
+          \ let g:vim_ai_is_selection_pending = mode() =~# "^[vV\<C-v>]"
+augroup END
+
+command! -range   -nargs=? AI        <line1>,<line2>call vim_ai#AIRun({}, <q-args>)
+command! -range   -nargs=? AIEdit    <line1>,<line2>call vim_ai#AIEditRun({}, <q-args>)
+" Whereas AI and AIEdit default to passing the current line as range
+" AIChat defaults to passing nothing which is achieved by -range=0 and passing
+" <count> as described at https://stackoverflow.com/a/20133772
+command! -range=0 -nargs=? AIChat    <line1>,<line2>call vim_ai#AIChatRun(<count>, {}, <q-args>)
+command! -nargs=?          AINewChat                call vim_ai#AINewChatRun(<q-args>)
+command!                   AIRedo                   call vim_ai#AIRedoRun()