check Python code snippets on the fly with pyrefly

pyrefly_operator.vim

" Check Python code snippets on the fly with pyrefly.
" Put into ~/.vim/plugins` and use the mapping `g4` on a visual selection or a text object.

if exists('g:loaded_pyrefly_operator') | finish | endif
let g:loaded_pyrefly_operator = 1

if !executable('pyrefly') | finish | endif

let s:pyrefly_efm = join([
      \ '%E%\s%#ERROR %m',
      \ '%W%\s%#WARN %m',
      \ '%N%\s%#NOTE %m',
      \ '%C%\s%#--> %f:%l:%c',
      \ '%-G%*\d\s%#\|%.%#',
      \ '%-G%\s%#\|%.%#',
      \ '%-G%\s%#',
      \ '%C[ \t] %.%#',
      \ '%-G%.%#',
      \ ], ',')

function! s:PyreflyCmd() abort
  return get(b:, 'pyrefly_makeprg', get(g:, 'pyrefly_makeprg', 'pyrefly'))
endfunction

function! s:SetDiagnostics(target, items) abort
  if empty(a:items)
    if a:target ==# 'qf'
      cclose
      call setqflist([], 'r')
    else
      lclose
      call setloclist(0, [], 'r')
    endif
    return
  endif
  if a:target ==# 'qf'
    call setqflist(a:items, 'r')
    copen
  else
    call setloclist(0, a:items, 'r')
    lopen
  endif
endfunction

function! PyreflyCheckRange(line1, line2, ...) abort
  let opts = a:0 ? a:1 : {}
  let target = get(opts, 'target', get(g:, 'pyrefly_target', 'loc'))
  let bufnr0 = get(opts, 'bufnr', bufnr('%'))
  let cmd0 = get(opts, 'cmd', s:PyreflyCmd())

  let l1 = a:line1
  let l2 = a:line2
  if l1 <= 0 || l2 <= 0
    return
  endif
  if l2 < l1
    let [l1, l2] = [l2, l1]
  endif

  let lines = getbufline(bufnr0, l1, l2)
  if empty(lines)
    call s:SetDiagnostics(target, [])
    return
  endif

  " Default: use pyrefly snippet on stdin.
  " Allow override via g:pyrefly_snippet_args.
  let snippet_args = get(g:, 'pyrefly_snippet_args', 'snippet --color=never --summary=none --output-format=full-text')
  let input = join(lines, "\n") .. "\n"
  let out = systemlist(cmd0 .. ' ' .. snippet_args .. ' ' .. shellescape(input))

  " Parse without modifying existing quickfix/location list.
  let title = printf('pyrefly: %s', bufname(bufnr0))
  let parsed = getqflist({'lines': out, 'efm': s:pyrefly_efm, 'title': title})
  let items = get(parsed, 'items', [])

  " Map snippet line numbers back to buffer line numbers.
  let offset = l1 - 1
  for item in items
    if get(item, 'lnum', 0) > 0
      let item.lnum += offset
    endif
    if has_key(item, 'end_lnum') && item.end_lnum > 0
      let item.end_lnum += offset
    endif

    let item.bufnr = bufnr0
    if has_key(item, 'filename')
      unlet item.filename
    endif
  endfor

  call s:SetDiagnostics(target, items)
endfunction

function! PyreflyOperator(type) abort
  " Called by g@ after motion; uses '[ and '] marks.
  let l1 = line("'[")
  let l2 = line("']")
  if l1 <= 0 || l2 <= 0
    return
  endif
  call PyreflyCheckRange(l1, l2, {'target': get(g:, 'pyrefly_target', 'loc')})
endfunction

" Normal-mode operator: type <Leader>p{motion}, e.g. <Leader>pip or <Leader>p2j.
nnoremap g4 <Cmd>let &operatorfunc = function('PyreflyOperator')<CR>g@
" Visual-mode mapping: run on selected lines directly.
xnoremap g4 <Cmd>call PyreflyCheckRange(line("'<"), line("'>"), {'target': 'loc'})<CR>