いろいろ備忘録日記

主に .NET とか Go とか Python絡みのメモを公開しています。

Goメモ-70 (チャネルから指定された個数分取得するチャネル, Take, TakeWhile, TakeWhileFn)

概要

チャネル関数の続き。

今回は、指定されたチャネルから指定された個数分だけデータを取り出すチャネルを返す関数です。

C# の Enumerable.TakeEnumerable.TakeWhile チックな感じ。

サンプル

package chans

// Take -- 指定した個数分、入力用チャネルから値を取得するチャネルを返します。
func Take(done <-chan struct{}, in <-chan interface{}, count int) <-chan interface{} {
    out := make(chan interface{})

    go func() {
        defer close(out)

        for i := 0; i < count; i++ {
            select {
            case <-done:
                return
            case out <- <-in:
            }
        }
    }()

    return out
}

// TakeWhile -- 入力用チャネルから取得した値が指定した値と同一である間、値を取得し続けるチャネルを返します。
func TakeWhile(done <-chan struct{}, in <-chan interface{}, value interface{}) <-chan interface{} {
    return TakeWhileFn(done, in, func() interface{} { return value })
}

// TakeWhileFn -- 入力用チャネルから取得した値が指定した関数の戻り値と同一である間、値を取得し続けるチャネルを返します。
func TakeWhileFn(done <-chan struct{}, in <-chan interface{}, fn func() interface{}) <-chan interface{} {
    out := make(chan interface{})

    go func(fn func() interface{}) {
        defer close(out)

        r := fn()
        for {
            select {
            case <-done:
                return
            case v, ok := <-in:
                if !ok {
                    return
                }

                if v != r {
                    return
                }

                select {
                case out <- v:
                case <-done:
                }
            }
        }
    }(fn)

    return out
}

gomy/take.go at master · devlights/gomy · GitHub

以下、テストコードです。

package chans

import (
    "testing"
)

func TestTake(t *testing.T) {
    type (
        testin struct {
            total int
            count int
        }
        testout struct {
            count int
        }
        testcase struct {
            in  testin
            out testout
        }
    )

    cases := []testcase{
        {
            in: testin{
                total: 10,
                count: 0,
            },
            out: testout{count: 0},
        },
        {
            in: testin{
                total: 10,
                count: 1,
            },
            out: testout{count: 1},
        },
        {
            in: testin{
                total: 10,
                count: 5,
            },
            out: testout{count: 5},
        },
        {
            in: testin{
                total: 10,
                count: 10,
            },
            out: testout{count: 10},
        },
    }

    for caseCount, c := range cases {
        func() {
            done := make(chan struct{})
            defer close(done)

            inCh := make(chan interface{}, c.in.total)
            func() {
                defer close(inCh)

                for i := 0; i < c.in.total; i++ {
                    inCh <- i
                }
            }()

            takeCh := Take(done, inCh, c.in.count)

            recvCount := 0
            for v := range takeCh {
                t.Logf("[test-%02d][take] %v\n", caseCount+1, v)
                recvCount++
            }

            if c.out.count != recvCount {
                t.Errorf("want: %v\tgot: %v", c.out.count, recvCount)
            }
        }()
    }
}

func TestTakeWhile(t *testing.T) {
    type (
        testin struct {
            value int
            data  []int
        }
        testout struct {
            count int
        }
        testcase struct {
            in  testin
            out testout
        }
    )

    cases := []testcase{
        {
            in: testin{
                data:  []int{1},
                value: 1,
            },
            out: testout{count: 1},
        },
        {
            in: testin{
                data:  []int{1, 1, 1, 1, 1, 2, 2},
                value: 1,
            },
            out: testout{count: 5},
        },
        {
            in: testin{
                data:  []int{1, 1, 2, 1, 1, 2, 2},
                value: 1,
            },
            out: testout{count: 2},
        },
    }

    for caseCount, c := range cases {
        func() {
            done := make(chan struct{})
            defer close(done)

            inCh := make(chan interface{}, len(c.in.data))
            func() {
                defer close(inCh)

                for _, v := range c.in.data {
                    inCh <- v
                }
            }()

            takeCh := TakeWhile(done, inCh, c.in.value)

            recvCount := 0
            for v := range takeCh {
                t.Logf("[test-%02d][take] %v\n", caseCount+1, v)
                recvCount++
            }

            if c.out.count != recvCount {
                t.Errorf("want: %v\tgot: %v", c.out.count, recvCount)
            }
        }()
    }
}

func TestTakeWhileFn(t *testing.T) {
    type (
        testin struct {
            fn   func() interface{}
            data []int
        }
        testout struct {
            count int
        }
        testcase struct {
            in  testin
            out testout
        }
    )

    cases := []testcase{
        {
            in: testin{
                data: []int{1},
                fn:   func() interface{} { return 1 },
            },
            out: testout{count: 1},
        },
        {
            in: testin{
                data: []int{1, 1, 1, 1, 1, 2, 2},
                fn:   func() interface{} { return 1 },
            },
            out: testout{count: 5},
        },
        {
            in: testin{
                data: []int{1, 1, 2, 1, 1, 2, 2},
                fn:   func() interface{} { return 1 },
            },
            out: testout{count: 2},
        },
    }

    for caseCount, c := range cases {
        func() {
            done := make(chan struct{})
            defer close(done)

            inCh := make(chan interface{}, len(c.in.data))
            func() {
                defer close(inCh)

                for _, v := range c.in.data {
                    inCh <- v
                }
            }()

            takeCh := TakeWhileFn(done, inCh, c.in.fn)

            recvCount := 0
            for v := range takeCh {
                t.Logf("[test-%02d][take] %v\n", caseCount+1, v)
                recvCount++
            }

            if c.out.count != recvCount {
                t.Errorf("want: %v\tgot: %v", c.out.count, recvCount)
            }
        }()
    }
}

実行すると以下のようになります。

$ go test -v github.com/devlights/gomy/chans -run ^TestTake.*$
=== RUN   TestTake
    TestTake: take_test.go:71: [test-02][take] 0
    TestTake: take_test.go:71: [test-03][take] 0
    TestTake: take_test.go:71: [test-03][take] 1
    TestTake: take_test.go:71: [test-03][take] 2
    TestTake: take_test.go:71: [test-03][take] 3
    TestTake: take_test.go:71: [test-03][take] 4
    TestTake: take_test.go:71: [test-04][take] 0
    TestTake: take_test.go:71: [test-04][take] 1
    TestTake: take_test.go:71: [test-04][take] 2
    TestTake: take_test.go:71: [test-04][take] 3
    TestTake: take_test.go:71: [test-04][take] 4
    TestTake: take_test.go:71: [test-04][take] 5
    TestTake: take_test.go:71: [test-04][take] 6
    TestTake: take_test.go:71: [test-04][take] 7
    TestTake: take_test.go:71: [test-04][take] 8
    TestTake: take_test.go:71: [test-04][take] 9
--- PASS: TestTake (0.00s)
=== RUN   TestTakeWhile
    TestTakeWhile: take_test.go:139: [test-01][take] 1
    TestTakeWhile: take_test.go:139: [test-02][take] 1
    TestTakeWhile: take_test.go:139: [test-02][take] 1
    TestTakeWhile: take_test.go:139: [test-02][take] 1
    TestTakeWhile: take_test.go:139: [test-02][take] 1
    TestTakeWhile: take_test.go:139: [test-02][take] 1
    TestTakeWhile: take_test.go:139: [test-03][take] 1
    TestTakeWhile: take_test.go:139: [test-03][take] 1
--- PASS: TestTakeWhile (0.00s)
=== RUN   TestTakeWhileFn
    TestTakeWhileFn: take_test.go:207: [test-01][take] 1
    TestTakeWhileFn: take_test.go:207: [test-02][take] 1
    TestTakeWhileFn: take_test.go:207: [test-02][take] 1
    TestTakeWhileFn: take_test.go:207: [test-02][take] 1
    TestTakeWhileFn: take_test.go:207: [test-02][take] 1
    TestTakeWhileFn: take_test.go:207: [test-02][take] 1
    TestTakeWhileFn: take_test.go:207: [test-03][take] 1
    TestTakeWhileFn: take_test.go:207: [test-03][take] 1
--- PASS: TestTakeWhileFn (0.00s)
PASS
ok      github.com/devlights/gomy/chans 0.014s

前の記事のRepeatと今回のTakeを合わせると以下のように無限シーケンスから必要な分だけ取得するといったことが出来たりします。

例えば、こんな感じです。(前回と同じサンプルです)

package async

import (
    "context"
    "fmt"
    "math/rand"

    "github.com/devlights/gomy/chans"
)

// TakeFirst10Items -- 最初の10個のみを取得するサンプルです
func TakeFirst10Items() error {
    var (
        rootCtx         = context.Background()
        mainCtx, cancel = context.WithCancel(rootCtx)
    )

    defer cancel()

    // 乱数を返す関数
    randomInt := func() interface{} {
        return rand.Int()
    }

    // 乱数を延々と返すチャネル生成
    repeatCh := chans.RepeatFn(mainCtx.Done(), randomInt)

    // 最初の10件のみ取得するチャネル生成
    takeFirst10ItemCh := chans.Take(mainCtx.Done(), repeatCh, 10)

    // 出力
    for v := range takeFirst10ItemCh {
        fmt.Println(v)
    }

    return nil
}

try-golang/take_first_10items.go at master · devlights/try-golang · GitHub

実行すると以下のようになります。

$ make run
ENTER EXAMPLE NAME: async_take_first_10items
[Name] "async_take_first_10items"
711512294376717777
6597590559487096282
7546803073246591619
2730330101447707632
8913116274994349004
1574111430622025378
963317749379199556
5718127538066206519
3934619028980328702
7294862805965956347


[Elapsed] 2.340763ms

参考

Go言語による並行処理

Go言語による並行処理

関連記事

devlights.hatenablog.com

devlights.hatenablog.com

devlights.hatenablog.com


過去の記事については、以下のページからご参照下さい。

  • いろいろ備忘録日記まとめ

devlights.github.io

サンプルコードは、以下の場所で公開しています。

  • いろいろ備忘録日記サンプルソース置き場

github.com

github.com

github.com