Golang's equivalent of itertools.chain?-CodePudding

In Golang, how can I iterate over three slices with a single for loop, without creating a new slice containing copies of the elements in all three slices? Is there something like Python's itertools.chain?

CodePudding user response：

You could use channels to build your own chain() method.

package main

import "fmt"

func main() {
    slice1 := []int{2, 3, 5, 7, 11, 13}
    slice2 := []int{21321, 12313, 213}
    slice3 := []int{8987, 988, 675676, 6587686}

    for val := range chain(slice1, slice2, slice3) {
        fmt.Println(val)
    }
}

func chain[T any](slices ...[]T) <-chan T {
    channel := make(chan T)
    go func() {
        for _, slice := range slices {
            for _, val := range slice {
                channel <- val
            }
        }
        close(channel)
    }()
    return channel
}

Here the link to Go Playground.

Essentially the idea is to use variadic argument slices which can hold and undetermined number of slices and in the function we then create a channel of the given type return that channel so a caller can retrieve values from it and before we do that we start a Goroutine which will actually send the values over the channel. In order to be able to use range we need to close() the channel once we have looped over all slices.

This uses Generics which are available as of Go 1.18 which has just recently been released.

Please note: I am myself quite new to Go, so there might be better ways to do this. Having said this, I am happy to include any helpful suggestions.

CodePudding user response：

A simple solution using generics

package main

import "fmt"

func Chain[T any](f func(e T), slices ...[]T) {
    for _, slice := range slices {
        for _, e := range slice {
            f(e)
        }
    }
}

func main() {
    slice1 := []int{1, 2, 3}
    slice2 := []int{10, 20, 30}
    slice3 := []int{100, 200, 300}

    Chain(func(e int) {
        fmt.Println(e)
    }, slice1, slice2, slice3)
}

Output:

The function Chain takes a function parameter, which will be executed for each element in each slice consecutively. This function will serve as your loop body code.

The solution can be extended to retain other looping features such as break and index numbers:

func Chain[T any](f func(i int, e T) bool, slices ...[]T) {
    var i int
    for _, slice := range slices {
        for _, e := range slice {
            if !f(i, e) {
                return
            }
            i  
        }
    }
}

...

    Chain(func(i int, e int) bool {
        fmt.Println(i, "-", e)
        return (e <= 20)
    }, slice1, slice2, slice3)
...

The loop will "break" when f returns false.

Output: