Moved sagetex syntax changes to respective file

autoload/vimtex/syntax.vim

@@ -24,19 +24,14 @@ function! vimtex#syntax#in_comment(...) abort " {{{1
   return call('vimtex#syntax#in', ['texComment'] + a:000)
 endfunction
 
-" }}}1
-function! vimtex#syntax#in_sage(...) abort " {{{1
-  return call('vimtex#syntax#in', ['texSagetexArg'] + a:000) || call('vimtex#syntax#in', ['texSagetexZone'] + a:000)
-endfunction
-
 " }}}1
 function! vimtex#syntax#in_mathzone(...) abort " {{{1
   " The following checks if we are inside a texMathZone environment. The
-  " arguments to \label{...}, the texRefArg group, the texSagetexArg group,
-  " and \text{...} like commands, the texMathTextArg group, are actively
-  " ignored as these should not be considered to be math environments.
+  " arguments to \label{...}, the texRefArg group, and \text{...} like
+  " commands, the texMathTextArg group, are actively ignored as these should
+  " not be considered to be math environments.
   let l:groups = reverse(call('vimtex#syntax#stack', a:000))
-  let l:group = matchstr(l:groups, '\v^tex%(Math%(Zone|Text|Tag)|RefArg|SagetexArg)')
+  let l:group = matchstr(l:groups, '\v^tex%(Math%(Zone|Text|Tag)|RefArg)')
   return l:group =~# '^texMathZone'
 endfunction
 

autoload/vimtex/syntax/p/sagetex.vim

@@ -16,34 +16,40 @@ function! vimtex#syntax#p#sagetex#load(cfg) abort " {{{1
         \})
   call vimtex#syntax#core#new_opt('texSagetexOpt', {'next': 'texSagetexArg'})
 
-  for l:env in [
+  for l:env_name in [
         \ 'sageblock',
         \ 'sagesilent',
         \ 'sageverbatim',
         \ 'sageexample',
         \ 'sagecommandline'
         \]
     call vimtex#syntax#core#new_env({
-          \ 'name': l:env,
+          \ 'name': l:env_name,
           \ 'region': 'texSagetexZone',
           \ 'contains': '@vimtex_nested_python'
           \})
   endfor
 
-  for l:env in [
+  " Note: The following commands are supported both inside and outside math zones
+  for l:cmd_name in [
         \ 'sage',
         \ 'sagestr'
         \]
-    execute 'syntax match texCmdSagetex /\\' . l:env .
-          \ '\>/ nextgroup=texSagetexArg skipwhite skipnl'
-    call vimtex#syntax#core#new_cmd({
-          \ 'name': l:env,
-          \ 'mathmode': 1,
-          \ 'nextgroup': 'texSagetexArg'
-          \})
+    for l:in_mathmode in [v:true, v:false]
+      call vimtex#syntax#core#new_cmd({
+            \ 'name': l:cmd_name,
+            \ 'mathmode': l:in_mathmode,
+            \ 'nextgroup': 'texSagetexArg'
+            \})
+    endfor
   endfor
 
   highlight def link texCmdSagetex texCmd
 endfunction
 
 " }}}1
+function! vimtex#syntax#p#sagetex#in_sage(...) abort " {{{1
+  return call('vimtex#syntax#in', ['texSagetex\(Arg\|Zone\)'] + a:000)
+endfunction
+
+" }}}1