Golang's equivalent of itertools.chain?

user_14416704 0
Viewed 145

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?

go
2 Answers

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:

1
2
3
10
20
30
100
200
300

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:

0 - 1
1 - 2
2 - 3
3 - 10
4 - 20
5 - 30
user_114246730

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.

Edit

As correctly indicated by Zombo you will leak a goroutine if you break in the range loop above as not all values sent to the channel will be retrieved from it.

In order to get a break behavior without leaking a go routine you could pass a predicate function which, when it returns true will close the channel and therefore prevent a leak.

Similarly you could pass callback functions for all kinds of stuff (e.g. filtering) according to your use case.

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 chainUntil(func(item int) bool {
        return item > 6
    }, slice1, slice2, slice3) {
        fmt.Println(val)
    }
}

func chainUntil[T any](predicate func(item T) bool, slices ...[]T) <-chan T {
    channel := make(chan T)
    go func() {
        for _, slice := range slices {
            for _, val := range slice {
                if predicate(val) == true {
                    close(channel)
                    return
                }
                channel <- val
            }
        }
        close(channel)
    }()
    return channel
}

Expected output:

2
3
5
user_174873480
Related