How to allow count *before* my custom operator?

Question

I'd like to define my own operator. Vim's builtin help has a nice little tutorial on it, where they show you can create your own operator with opfunc and g@:

nmap <silent> <F4> :set opfunc=CountSpaces<CR>g@
vmap <silent> <F4> :<C-U>call CountSpaces(visualmode(), 1)<CR>

function! CountSpaces(type, ...)
  ...
endfunction

There's only one downside to this approach. Let's say for example, that I want to run <F4>2j. That works fine, but 2<F4>j does not. This gives

E481: No range allowed

This makes perfect sense because the range is applied to the set opfunc command, not the g@ command.

I thought that I would be able to get around it by doing this:

nnoremap <expr> <F4> ":\<C-u>set opfunc=Test\<cr>".v:count1.'g@'

But when I run this, it causes some strange issues where '[ and '] are incorrectly defined in the function. (I'd be happy to go into more details in chat, but that's unrelated to my question. I've reproduced this bug in Vim 8.0 on Windows and Ubuntu and in Neovim 0.2.1 on Windows)

So how can I define my operator so that a count works the same before or after the operator? I realize that this is a very minor difference, and that for most intents and purposes, I can simply get around it by always doing the count afterwards. But vim's default operators work this way, and I would really like mine too also.

Answer 1

You can call a function transparently within a mapping, thus circumventing mode changes (e.g. by pressing :) or losing counts, by using <expr> mappings, similar to what you suggested. Just use them to call a function like so:

function! SetOpFunc()
    set opfunc=CountSpaces
    return 'g@'
endfun
nnoremap <expr> <F4> SetOpFunc()

This technique isn't limited to operator mappings either. For example I've used it in Operator-pending mappings to call a function during the mapping without disrupting modes or counts, returning '' at the end. This way you can leverage Vim's native features rather than having to recreate them yourself.

Answer 2

Antony's answer is great, and I'd just like to add on some quality-of-life improvements.

With Antony's version, the count is not yet processed as of SetOpFunc, and is passed down to g@, so that the count passed to the custom command is multiplied with any counts in the following motion. For example, 2<custom-command>2E becomes <custom-command>4E.

If this is unwanted behaviour and you'd like to handle the count for the command separately from counts to the motion, you can use this solution.

Here, we negate the count passed to g@ by a black hole macro @_ that does nothing but eat the count, and handle the counts to your command manually by storing it. In your custom command, use the variables g:opcount and g:opcount1 instead of the usual v:count variables.

Also note that we've also made our SetOpFunc generic on the command function.

function! SetOpFunc(fn)
    execute 'set opfunc=' . a:fn
    let g:opcount = v:count
    let g:opcount1 = v:count1
    return '@_g@'
endfunction
nnoremap <expr> <F4> SetOpFunc('CountSpaces')