fix: cite context on \*cites{x}{y}

refer: lervag#2016

autoload/vimtex/cmd.vim

@@ -623,22 +623,25 @@ function! s:get_cmd(direction) abort " {{{1
     let res.pos_end.cnum = res.overlay.close.cnum
   endif
 
-  " Get options
-  let opt = s:get_cmd_part('[', res.pos_end)
-  while !empty(opt)
-    call add(res.opts, opt)
-    let res.pos_end.lnum = opt.close.lnum
-    let res.pos_end.cnum = opt.close.cnum
+  " Get options and arguments
+  while v:true
     let opt = s:get_cmd_part('[', res.pos_end)
-  endwhile
+    if !empty(opt)
+      call add(res.opts, opt)
+      let res.pos_end.lnum = opt.close.lnum
+      let res.pos_end.cnum = opt.close.cnum
+      continue
+    endif
 
-  " Get arguments
-  let arg = s:get_cmd_part('{', res.pos_end)
-  while !empty(arg)
-    call add(res.args, arg)
-    let res.pos_end.lnum = arg.close.lnum
-    let res.pos_end.cnum = arg.close.cnum
     let arg = s:get_cmd_part('{', res.pos_end)
+    if !empty(arg)
+      call add(res.args, arg)
+      let res.pos_end.lnum = arg.close.lnum
+      let res.pos_end.cnum = arg.close.cnum
+      continue
+    endif
+
+    break
   endwhile
 
   " Include entire cmd text

autoload/vimtex/context.vim

@@ -39,3 +39,30 @@ function! vimtex#context#menu() abort " {{{1
 endfunction
 
 " }}}1
+function! vimtex#context#get(...) abort " {{{1
+  if a:0 > 0
+    let l:pos_saved = vimtex#pos#get_cursor()
+    call vimtex#pos#set_cursor(a:000)
+  endif
+
+  let l:cmd = vimtex#cmd#get_current()
+  let l:word = expand('<cword>')
+
+  if a:0 > 0
+    call vimtex#pos#set_cursor(l:pos_saved)
+  endif
+
+  if empty(l:cmd) | return | endif
+
+  for l:handler in b:vimtex.context_menu
+    if l:handler.match(l:cmd, l:word)
+      return {
+            \ 'cmd': l:cmd,
+            \ 'word': l:word,
+            \ 'handler': l:handler,
+            \}
+    endif
+  endfor
+endfunction
+
+" }}}1

autoload/vimtex/context/cite.vim

@@ -19,11 +19,16 @@ function! s:handler.match(cmd, word) abort dict " {{{1
     return v:false
   endif
 
-  if len(a:cmd.args) != 1
+  if len(a:cmd.args) < 1 || len(a:cmd.args) > 2
     return v:false
   endif
 
-  let self.cites = split(a:cmd.args[0].text, ',\s*')
+  let l:text = a:cmd.args[0].text
+  if len(a:cmd.args) == 2
+    let l:text .= ',' . a:cmd.args[1].text
+  endif
+
+  let self.cites = split(l:text, ',\s*')
   if index(self.cites, a:word) >= 0
     let self.selected = a:word
   else

test/test-context-cite/Makefile

@@ -0,0 +1,14 @@
+MYVIM ?= nvim --headless
+
+INMAKE := 1
+export INMAKE
+
+TESTS := $(wildcard test*.vim)
+TESTS := $(TESTS:.vim=)
+
+.PHONY: test $(TESTS)
+
+test: $(TESTS)
+
+$(TESTS):
+	@$(MYVIM) -u $@.vim

test/test-context-cite/test-cites.bib

@@ -0,0 +1,13 @@
+@book{Hemingway1940,
+  author    = {Earnest Hemingway},
+  title     = {For Whom the Bell Tolls},
+  date      = {1940},
+  publisher = {Simon \& Schuster},
+}
+
+@book{JiM2020,
+  author    = {{Justice in Mexico}},
+  title     = {Organized Crime and Violence in Mexico: 2020 Special Report},
+  date      = {2019-07},
+  publisher = {Justice in Mexico},
+}

test/test-context-cite/test-cites.tex

@@ -0,0 +1,13 @@
+\documentclass{article}
+\usepackage{biblatex}
+\addbibresource{test-cites.bib}
+
+\begin{document}
+\begin{enumerate}
+  \item \autocite[10]{Hemingway1940}
+  \item \autocites{Hemingway1940}{JiM2020}
+  \item \autocites[10]{JiM2020}[20]{Hemingway1940}
+  \item \autocites[10]{Hemingway1940}{JiM2020}
+  \item \autocites{Hemingway1940}[20]{JiM2020}
+\end{enumerate}
+\end{document}

test/test-context-cite/test-cites.vim

@@ -0,0 +1,19 @@
+set nocompatible
+let &rtp = '../..,' . &rtp
+filetype plugin indent on
+syntax enable
+
+set nomore
+
+nnoremap q :qall!<cr>
+
+silent edit test-cites.tex
+
+if empty($INMAKE) | finish | endif
+
+call vimtex#test#assert_equal('Hemingway1940',
+      \ vimtex#context#get(9, 49).handler.selected)
+call vimtex#test#assert_equal('JiM2020',
+      \ vimtex#context#get(11, 39).handler.selected)
+
+quit!