Julia splat operator unpacking

In Python, one can use the * operator in the unpacking of an iterable.

In [1]: head, *tail = [1, 2, 3, 4, 5]

In [2]: head
Out[2]: 1

In [3]: tail
Out[3]: [2, 3, 4, 5]

I would like to produce the same behavior in Julia. I figured that the equivalent ... operator would work, but it seems to just produce an error in this context.

julia> head, tail... = [1, 2, 3, 4, 5]
ERROR: syntax: invalid assignment location "tail..."

I was able to produce the results I want using the following, but this is an ugly solution.

julia> head, tail = A[1], A[2:end]
(1,[2,3,4,5])

Can I unpack the array such that tail would contain the rest of the items after head using the splat (...) operator? If not, what is the cleanest alternative?

Edit: This feature has been proposed in #2626. It looks like it will be part of the 1.0 release.

Solution

That does indeed sound like a job for a macro:

function unpack(lhs, rhs) 
    len = length(lhs.args)
    if len == 1
        # just remove the splatting
        l, is_splat = remove_splat(lhs.args[1])
        return :($l = $(esc(rhs)))
    else
        new_lhs = :()
        new_rhs = quote 
            tmp = $(esc(rhs))
            $(Expr(:tuple)) 
        end
        splatted = false
        for (i, e) in enumerate(lhs.args)
            l, is_splat = remove_splat(e)
            if is_splat
                splatted && error("Only one splatting operation allowed on lhs")
                splatted = true
                r = :(tmp[$i:end-$(len-i)])
            elseif splatted
                r = :(tmp[end-$(len-i)])
            else
                r = :(tmp[$i])
            end
            push!(new_lhs.args, l)
            push!(new_rhs.args[4].args, r)
        end
        return :($new_lhs =  $new_rhs)
    end
end

remove_splat(e::Symbol) = esc(e),  false

function remove_splat(e::Expr)
    if e.head == :(...)
        return esc(e.args[1]), true
    else
        return esc(e), false
    end
end

macro unpack(expr)
    if Meta.isexpr(expr, :(=))
        if Meta.isexpr(expr.args[1], :tuple)
            return unpack(expr.args[1], expr.args[2])
        else
            return unpack(:(($(expr.args[1]),)), expr.args[2])
        end
    else
        error("Cannot parse expression")
    end
end

It is not very well tested, but basic things work:

julia> @unpack head, tail... = [1,2,3,4]
(1,[2,3,4])

julia> @unpack head, middle..., tail = [1,2,3,4,5]
(1,[2,3,4],5)

A few Julia gotchas:

x,y = [1,2,3] #=> x = 1, y = 2

a = rand(3)
a[1:3], y = [1,2,3] #=> a = [1.0,1.0,1.0], y = 2

The macro follows this behavior

@unpack a[1:3], y... = [1,2,3]
#=> a=[1.0,1.0,1.0], y=[2,3]