func zip()

in starlark/library.go [1138:1189]


func zip(thread *Thread, _ *Builtin, args Tuple, kwargs []Tuple) (Value, error) {
	if len(kwargs) > 0 {
		return nil, fmt.Errorf("zip does not accept keyword arguments")
	}
	rows, cols := 0, len(args)
	iters := make([]Iterator, cols)
	defer func() {
		for _, iter := range iters {
			if iter != nil {
				iter.Done()
			}
		}
	}()
	for i, seq := range args {
		it := Iterate(seq)
		if it == nil {
			return nil, fmt.Errorf("zip: argument #%d is not iterable: %s", i+1, seq.Type())
		}
		iters[i] = it
		n := Len(seq)
		if i == 0 || n < rows {
			rows = n // possibly -1
		}
	}
	var result []Value
	if rows >= 0 {
		// length known
		result = make([]Value, rows)
		array := make(Tuple, cols*rows) // allocate a single backing array
		for i := 0; i < rows; i++ {
			tuple := array[:cols:cols]
			array = array[cols:]
			for j, iter := range iters {
				iter.Next(&tuple[j])
			}
			result[i] = tuple
		}
	} else {
		// length not known
	outer:
		for {
			tuple := make(Tuple, cols)
			for i, iter := range iters {
				if !iter.Next(&tuple[i]) {
					break outer
				}
			}
			result = append(result, tuple)
		}
	}
	return NewList(result), nil
}