From f66c0938e5e35c46697654263f1c1c33482cddfa Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Fri, 15 Apr 2022 16:22:19 +0800
Subject: [PATCH] feat: add Take function to generate a chan of values taken
 from another chan

---
 concurrency/channel.go      | 19 +++++++++++++++
 concurrency/channel_test.go | 46 +++++++++++++++++++++++++++++--------
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/concurrency/channel.go b/concurrency/channel.go
index e45408f..d9d592f 100644
--- a/concurrency/channel.go
+++ b/concurrency/channel.go
@@ -52,3 +52,22 @@ func (c *Channel) Repeat(done <-chan any, values ...any) <-chan any {
 	}()
 	return dataStream
 }
+
+// Take return a chan whose values are tahken from another chan
+func (c *Channel) Take(done <-chan any, valueStream <-chan any, number int) <-chan any {
+	takeStream := make(chan any)
+
+	go func() {
+		defer close(takeStream)
+
+		for i := 0; i < number; i++ {
+			select {
+			case <-done:
+				return
+			case takeStream <- <-valueStream:
+			}
+		}
+	}()
+
+	return takeStream
+}
diff --git a/concurrency/channel_test.go b/concurrency/channel_test.go
index 2f1334e..6ee33bd 100644
--- a/concurrency/channel_test.go
+++ b/concurrency/channel_test.go
@@ -16,7 +16,7 @@ func TestGenerate(t *testing.T) {
 	intStream := c.Generate(done, 1, 2, 3)
 
 	// for v := range intStream {
-	// 	t.Log(v)
+	// 	t.Log(v) //1, 2, 3
 	// }
 	assert.Equal(1, <-intStream)
 	assert.Equal(2, <-intStream)
@@ -24,18 +24,46 @@ func TestGenerate(t *testing.T) {
 }
 
 func TestRepeat(t *testing.T) {
-	// assert := internal.NewAssert(t, "TestRepeat")
+	assert := internal.NewAssert(t, "TestRepeat")
 
 	done := make(chan any)
 	defer close(done)
 
 	c := NewChannel()
-	intStream := c.Repeat(done, 1, 2)
+	intStream := c.Take(done, c.Repeat(done, 1, 2), 5)
 
-	for v := range intStream {
-		t.Log(v)
-	}
-	// assert.Equal(1, <-intStream)
-	// assert.Equal(2, <-intStream)
-	// assert.Equal(3, <-intStream)
+	// for v := range intStream {
+	// 	t.Log(v) //1, 2, 1, 2, 1
+	// }
+	assert.Equal(1, <-intStream)
+	assert.Equal(2, <-intStream)
+	assert.Equal(1, <-intStream)
+	assert.Equal(2, <-intStream)
+	assert.Equal(1, <-intStream)
+}
+
+func TestTake(t *testing.T) {
+	assert := internal.NewAssert(t, "TestRepeat")
+
+	done := make(chan any)
+	defer close(done)
+
+	numbers := make(chan any, 5)
+	numbers <- 1
+	numbers <- 2
+	numbers <- 3
+	numbers <- 4
+	numbers <- 5
+	defer close(numbers)
+
+	c := NewChannel()
+	intStream := c.Take(done, numbers, 3)
+
+	// for v := range intStream {
+	// 	t.Log(v) //1, 2, 3
+	// }
+
+	assert.Equal(1, <-intStream)
+	assert.Equal(2, <-intStream)
+	assert.Equal(3, <-intStream)
 }