Wrap a python expression into a list comprehension

iter.vim

" Save as: ~/.vim/ftplugin/python/iter.vim
"
" Execute the command :Iter <type> to take the word under the cursor and turn
" it into a list comprehension. Example with "target" as the word:
"
" - :Iter set  => {v for v in target}
" - :Iter list => [v for v in target]
" - :Iter dict => {k: v for k, v in target.items()}
"
" If the command is called in visual mode, it will take the visual selection
" as the target rather than the word under the cursor.
"
scriptversion 4

let s:completion_types = ['dict', 'list', 'set']

command!
      \ -nargs=*
      \ -range=0
      \ -complete=custom,s:IterComplete
      \ Iter call s:Iter(<count>, <q-args>)

function! s:Iter(count, type) abort
  let saved_view = winsaveview()
  defer winrestview(saved_view)

  let count = a:count
  let type = a:type
  if type == ''
    let type = 'list'
  endif

  if count > 0
    let collection = trim(s:GetMotion('gv'))
  else
    let collection = expand('<cword>')
  endif

  if type == 'dict'
    if stridx(collection, ' ') >= 0 && collection !~ '[)}\]]$'
      " There's a space in the expression and it doesn't end in a bracket,
      " wrap it in ():
      let collection = $"({collection}).items()"
    else
      let collection = $"{collection}.items()"
    endif

    let comprehension = $"{{k: v for k, v in {collection}}}"
  elseif type == 'list'
    let comprehension = $"[v for v in {collection}]"
  elseif type == 'set'
    let comprehension = $"{{v for v in {collection}}}"
  else
    echoerr $"Type {type} not in {s:completion_types}"
    return
  endif

  if count > 0
    call s:ReplaceMotion('gv', comprehension)
  else
    call s:ReplaceMotion('viw', comprehension)
  endif
endfunction

function! s:IterComplete(argument_lead, command_line, cursor_position)
  return join(s:completion_types, "\n")
endfunction

function s:GetMotion(motion)
  let saved_register = getreg('a')
  defer setreg('a', saved_register)

  exe 'normal! ' .. a:motion .. '"ay'
  return @a
endfunction

function s:ReplaceMotion(motion, text)
  let saved_register = getreg('a')
  defer setreg('a', saved_register)

  let @a = a:text

  exe 'normal! ' .. a:motion .. '"ap'
endfunction