From 82cbb547879bb3ab5ed60bcbf57861a984c8e656 Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Tue, 17 Jan 2023 16:47:20 +0800
Subject: [PATCH] feat: add FromChannel for create stream

---
 stream/stream.go      | 17 ++++++++++++++---
 stream/stream_test.go | 16 ++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/stream/stream.go b/stream/stream.go
index 9bf4537..4ac92b4 100644
--- a/stream/stream.go
+++ b/stream/stream.go
@@ -57,8 +57,8 @@ func Of[T any](elems ...T) stream[T] {
 }
 
 // Generate stream where each element is generated by the provided generater function
-// generater function: func(args ...U) func() (item T, ok bool) {}
-func Generate[T any](generator func() func() (T, bool)) stream[T] {
+// generater function: func() func() (item T, ok bool) {}
+func Generate[T any](generator func() func() (item T, ok bool)) stream[T] {
 	source := make([]T, 0)
 
 	var zeroValue T
@@ -77,6 +77,17 @@ func FromSlice[T any](source []T) stream[T] {
 	return stream[T]{source: source}
 }
 
+// FromChannel create stream from channel.
+func FromChannel[T any](source <-chan T) stream[T] {
+	s := make([]T, 0)
+
+	for v := range source {
+		s = append(s, v)
+	}
+
+	return FromSlice(s)
+}
+
 // FromRange create a number stream from start to end. both start and end are included. [start, end]
 func FromRange[T constraints.Integer | constraints.Float](start, end, step T) stream[T] {
 	if end < start {
@@ -92,7 +103,7 @@ func FromRange[T constraints.Integer | constraints.Float](start, end, step T) st
 		source[i] = start + (T(i) * step)
 	}
 
-	return stream[T]{source: source}
+	return FromSlice(source)
 }
 
 // Distinct returns a stream that removes the duplicated items.
diff --git a/stream/stream_test.go b/stream/stream_test.go
index 0cc41be..271e46f 100644
--- a/stream/stream_test.go
+++ b/stream/stream_test.go
@@ -39,6 +39,22 @@ func TestFromSlice(t *testing.T) {
 	assert.Equal([]int{1, 2, 3}, stream.ToSlice())
 }
 
+func TestFromChannel(t *testing.T) {
+	assert := internal.NewAssert(t, "TestFromChannel")
+
+	ch := make(chan int)
+	go func() {
+		for i := 1; i < 4; i++ {
+			ch <- i
+		}
+		close(ch)
+	}()
+
+	stream := FromChannel(ch)
+
+	assert.Equal([]int{1, 2, 3}, stream.ToSlice())
+}
+
 func TestFromRange(t *testing.T) {
 	assert := internal.NewAssert(t, "TestFromRange")