take

2 files changed, 43 insertions, 2 deletions

diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam
index 53b80d3..6634f23 100644
--- a/src/gleam/iterator.gleam
+++ b/src/gleam/iterator.gleam
@@ -63,7 +63,6 @@ pub fn repeat(x: element) -> Iterator(element) {
   repeatedly(fn() { x })
 }
 
-// TODO: test
 // TODO: document
 pub fn from_list(list: List(element)) -> Iterator(element) {
   unfold(list, fn(acc) {
@@ -96,7 +95,6 @@ pub fn run(iterator) -> Nil {
   fold(iterator, Nil, fn(_, acc) { acc })
 }
 
-// TODO: test
 // TODO: document
 pub fn to_list(iterator: Iterator(element)) -> List(element) {
   iterator
@@ -104,6 +102,24 @@ pub fn to_list(iterator: Iterator(element)) -> List(element) {
   |> list.reverse
 }
 
+fn do_take(iterator, desired, acc) {
+  case desired > 0 {
+    True -> case iterator() {
+      Continue(element, iterator) -> do_take(iterator, desired - 1, [element | acc])
+      Stop -> acc
+    }
+    False -> acc
+  }
+}
+
+// TODO: document
+pub fn take(from iterator: Iterator(e), up_to desired: Int) -> List(e) {
+  iterator
+  |> unopaque
+  |> do_take(_, desired, [])
+  |> list.reverse
+}
+
 // Transforming Iterators
 
 fn do_map(iterator, f) {
diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam
index fbbf1c6..2d9ee3f 100644
--- a/test/gleam/iterator_test.gleam
+++ b/test/gleam/iterator_test.gleam
@@ -1,6 +1,10 @@
 import gleam/expect
 import gleam/iterator
+import gleam/list
 
+// TODO: Property tests
+
+// a |> from_list |> to_list == a
 pub fn to_from_list_test() {
   let test = fn(subject) {
     subject
@@ -14,3 +18,24 @@ pub fn to_from_list_test() {
   test([1, 2])
   test([1, 2, 4, 8])
 }
+
+// a |> from_list |> take(_, n) == a |> list.take(_, n)
+pub fn take_test() {
+  let test = fn(n, subject) {
+    subject
+    |> iterator.from_list
+    |> iterator.take(_, n)
+    |> expect.equal(_, list.take(subject, n))
+  }
+
+  test(0, [])
+  test(1, [])
+  test(-1, [])
+  test(0, [0])
+  test(1, [0])
+  test(-1, [0])
+  test(0, [0, 1, 2, 3, 4])
+  test(1, [0, 1, 2, 3, 4])
+  test(2, [0, 1, 2, 3, 4])
+  test(22, [0, 1, 2, 3, 4])
+}