How to make my python zip but in Lua work?

I am working on a library extension module for roblox. I really like Python’s zip() and I thought I should implement it. I have never written my own iterator before so I probably return something wrong.

zip = function(_t1, _t2)
	return function(t1, t2)
		return select(2, next(t1)), select(2, next(t2))
	end, _t1, _t2
end

It seems to work fine the first iteration

for v1, v2 in zip({ 0, 1, 2 }, { 3, 4, 5 }) do
    print(v1)
    print(v2)
end

It prints 0 but after that an exception bad argument #1 (table expected, got number) And I don’t see how a number was passed.

I added print(type(t1), type(t2)) before return select(2, next(t1)), select(2, next(t2)) to debug and I got this output

table	table
0
table	number
bad argument #1 (table expected, got number)
1 Like

next(t1) and next(t2) will always return the first argument of t1 and t2. next(t, k) returns the keyvalue-pair of t after the one with key k. If k is nil, as in your example, it returns the first keyvalue pair.

Here’s my implementation of a zip iterator. I’m not 100% sure it’s actually the same as the standard python one tho.

local t1, t2 = {1, 2, 3, 4}, {4, 5, 6, 7, 8}
local t3, t4 = {'a', 'b', 'c'}, {1, 2}

function zip(t1, t2, i)
    lt1, lt2 = #t1, #t2
    return function()
        i = (i or 0) + 1
        if i <= lt1 and i <= lt2  then
            return t1[i], t2[i]
        end
    end
end

for a, b in zip(t1, t2) do
    print(a, b)
end

for a, b in zip(t3, t4) do
    print(a, b)
end

Output:

line a b
1 1 4
2 2 5
3 3 6
3 4 7
4 a 1
5 b 2

EDIT: I think precomputing #t1 and #t2 might actually slow things down, because it creates a new closure. Can someone confirm/deny this?