" Python imports insert completion
" Maintainer: Sergey Matveev <stargrave@stargrave.org>
" License: GNU General Public License version 3 of the License or later
"
" This plugin is intended for quick import-line insertion.
" For example you have got "from foo.bar import Baz" import somewhere in
" your Git-versioned code. In another file you can write "Baz" and then
" press <F3> in insert mode. Completion menu will show you suggestions
" of all import lines from your code containing Baz import.
"
" If you have Pylint or Pyflakes output in quickfix window, containing
" unfedined variables errors, then you can you :call AllImportCompl()
" function. It will go through the quickfix list and find all possible
" imports, sort them and insert under current cursor position.
"
" It uses external "git grep" call and assumes that "canonical" Python
" import format is used (single import per line).
" You can use http://git.stargrave.org/cgit.cgi/pyimportcan.git/tree/pyimportcan.pl
" utility to convert existing imports to that format.

let s:git_grep_cmd = "git grep -H --line-number --ignore-case --no-color "

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

function! SortByLen(s1, s2)
    if len(a:s1) == len(a:s2) | return a:s1 > a:s2 | endif
    return 1 ? len(a:s1) > len(a:s2) : -1
endfunction

function! ImportCompl()
    normal diw
    let output = system(s:git_grep_cmd . '"^from .* import .*' . @" . '" -- "*.py" "**/*.py"')
    let suggestions = []
    for line in split(output, "\n")
        if stridx(line, "unused-import") != -1 | continue | endif
        let m = matchlist(line, '^.*:\d\+:\(.*\)$')
        if len(m) == 0 | continue | endif
        call insert(suggestions, m[1])
    endfor
    call sort(suggestions, "SortByLen")
    call uniq(suggestions)
    call reverse(suggestions)
    call complete(col('.'), suggestions)
    return ''
endfunction

inoremap <F3> <C-R>=ImportCompl()<CR>

function! AllImportCompl()
    let output = system(s:git_grep_cmd . '"^from .* import" -- "*.py" "**/*.py"')
    let imports = {}
    for line in split(output, "\n")
        if stridx(line, "unused-import") != -1 | continue | endif
        for regexp in [
            \'^.*:\d\+:\(from .* import \(\w\+\).*\)$',
            \'^.*:\d\+:\(from .* import \w\+ as \(\w\+\).*\)$',
        \]
            let m = matchlist(line, regexp)
            if len(m) == 0 | break | endif
            let imports[m[2]] = m[1]
        endfor
    endfor
    let lines = getloclist(winnr())
    if len(lines) == 0 | let lines = getqflist() | endif
    let result = []
    for line in lines
        let m = matchlist(line.text, '\(E0602\|F821\).*' . "'" . '\(\w\+\)' . "'$")
        if len(m) == 0 || !has_key(imports, m[2]) | continue | endif
        call insert(result, imports[m[2]])
    endfor
    call sort(result, "i")
    call uniq(result)
    call append(".", result)
endfunction