vim9script

# Function call splitter
# Maintainer: Sergey Matveev <stargrave@stargrave.org>
# License: GNU General Public License version 3 of the License or later
#
# This plugin splits function call on several lines.
#
#   def foobar(self, foo: str, bar: Some[thing, too]) -> None:
# to
#   def foobar(
#           self,
#           foo: str,
#           bar: Some[thing, too],
#   ) -> None:
#
#   foo(bar, baz)[0]
# to
#   foo(
#       bar,
#       baz,
#   )[0]
#
# You can un-split it using :Undefsplit command on the line where
# splitting starts.
#
# :Defsplit has optional argument specifying how many opening round
# parenthesis must be skipped.
# :Defsplit 1 on foo(baz(baz(...))) produces
#    foo(baz(
#        baz(...),
#    ))
#
# Also there is :Brsplit command behaving similarly, but it splits other
# types of brackets: "{}", "[]".

def BracketFind(brsAllowable: list<string>, line: string, offset: number): number
    var possible: list<number>
    var found: number
    for bracket in brsAllowable
        found = stridx(line, bracket, offset)
        if found != -1 | possible += [found] | endif
    endfor
    return min(possible)
enddef

const Brs = {"(": ")", "[": "]", "{": "}"}

export def Do(brsAllowable: list<string>, singleLineComma: bool, ...args: list<string>)
    var skip = len(args) == 0 ? 0 : str2nr(args[0])
    var shift = get(b:, "defsplit_shift", "    ")
    var line = getline(".")
    var prfx: string
    for i in range(len(line))
        if line[i] != shift[0]
            prfx = strpart(line, 0, i)
            if line[i : i + 3] ==# "def " ||
                line[i : i + 5] ==# "class " ||
                line[i : i + 9] ==# "async def "
                shift ..= shift
            endif
            break
        endif
    endfor
    var brfirst = BracketFind(brsAllowable, line, 0)
    var brlast = strridx(line, Brs[line[brfirst]])
    while skip > 0
        brfirst = BracketFind(brsAllowable, line, brfirst + 1)
        brlast = strridx(line, Brs[line[brfirst]], brlast - 1)
        skip -= 1
    endwhile
    var [curly, round, squar, outbuf] = [0, 0, 0, ""]
    var ready = [strpart(line, 0, brfirst + 1)]
    var trailingComma = true
    for c in split(line[brfirst + 1 : brlast - 1], '\zs')
        if c ==# "*" | trailingComma = false | endif
        if outbuf ==# "" && c ==# " " | continue | endif
        outbuf ..= c
        if c ==# "," && !curly && !round && !squar
            ready = add(ready, prfx .. shift .. outbuf)
            outbuf = ""
        elseif c ==# "[" | squar += 1
        elseif c ==# "]" | squar -= 1
        elseif c ==# "(" | round += 1
        elseif c ==# ")" | round -= 1
        elseif c ==# "{" | curly += 1
        elseif c ==# "}" | curly -= 1
        endif
    endfor
    if trailingComma && !(singleLineComma == true && len(ready) == 1)
        outbuf = outbuf .. ","
    endif
    ready = add(ready, prfx .. shift .. outbuf)
    ready = add(ready, prfx .. strpart(line, brlast))
    append(line("."), ready)
    normal "_dd
enddef