From aee6cfb82dffdddb54ab987731b3247721fa88c8 Mon Sep 17 00:00:00 2001
From: Luciano Ramalho <luciano@ramalho.org>
Date: Mon, 8 Jun 2026 23:43:40 -0300
Subject: [PATCH] start itertools package

---
 itertools/README.md     |  4 ++++
 itertools/basic.go      | 43 +++++++++++++++++++++++++++++++++
 itertools/basic_test.go | 53 +++++++++++++++++++++++++++++++++++++++++
 itertools/go.mod        |  3 +++
 4 files changed, 103 insertions(+)
 create mode 100644 itertools/README.md
 create mode 100644 itertools/basic.go
 create mode 100644 itertools/basic_test.go
 create mode 100644 itertools/go.mod

diff --git a/itertools/README.md b/itertools/README.md
new file mode 100644
index 0000000..b81e673
--- /dev/null
+++ b/itertools/README.md
@@ -0,0 +1,4 @@
+# itertools
+
+A collection of Iterator functions inspired by Python's `itertools` module,
+adapted to the conventions of Go.
diff --git a/itertools/basic.go b/itertools/basic.go
new file mode 100644
index 0000000..b96c71d
--- /dev/null
+++ b/itertools/basic.go
@@ -0,0 +1,43 @@
+package itertools
+
+import "iter"
+
+// Take returns a lazy iterator over the first n values of seq. It stops seq
+// early once n values have been yielded, so it is safe to apply to infinite
+// sequences like Count. Use slices.Collect to materialize the result.
+func Take(n int, seq iter.Seq[int]) iter.Seq[int] {
+	return func(yield func(int) bool) {
+		if n <= 0 {
+			return
+		}
+		taken := 0
+		for v := range seq {
+			if !yield(v) {
+				return
+			}
+			taken++
+			if taken == n {
+				return
+			}
+		}
+	}
+}
+
+// Count returns an infinite iterator that yields start, start+1, start+2, …
+// It is equivalent to Python's itertools.count(start).
+func Count(start int) iter.Seq[int] {
+	return CountStep(start, 1)
+}
+
+// CountStep returns an infinite iterator that yields start, start+step,
+// start+2*step, … — the Go equivalent of Python's itertools.count(start, step).
+// The caller must break or use a bounded helper (e.g. Take) to stop it.
+func CountStep(start, step int) iter.Seq[int] {
+	return func(yield func(int) bool) {
+		for i := start; ; i += step {
+			if !yield(i) {
+				return
+			}
+		}
+	}
+}
diff --git a/itertools/basic_test.go b/itertools/basic_test.go
new file mode 100644
index 0000000..101c7e1
--- /dev/null
+++ b/itertools/basic_test.go
@@ -0,0 +1,53 @@
+package itertools
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestCount(t *testing.T) {
+	got := slices.Collect(Take(5, Count(3)))
+	want := []int{3, 4, 5, 6, 7}
+	if !slices.Equal(got, want) {
+		t.Errorf("got %v, want %v", got, want)
+	}
+}
+
+func TestCountStep(t *testing.T) {
+	cases := []struct {
+		name  string
+		start int
+		step  int
+		len   int
+		want  []int
+	}{
+		{"step 1", 0, 1, 5, []int{0, 1, 2, 3, 4}},
+		{"non-zero start", 10, 1, 4, []int{10, 11, 12, 13}},
+		{"step > 1", 0, 10, 5, []int{0, 10, 20, 30, 40}},
+		{"negative step", 10, -3, 4, []int{10, 7, 4, 1}},
+		{"negative start", -5, 2, 4, []int{-5, -3, -1, 1}},
+		{"take zero", 0, 1, 0, nil},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			got := slices.Collect(Take(tc.len, CountStep(tc.start, tc.step)))
+			if !slices.Equal(got, tc.want) {
+				t.Errorf("got %v, want %v", got, tc.want)
+			}
+		})
+	}
+}
+
+func TestCountStep_earlyBreak(t *testing.T) {
+	// Verify the iterator stops cleanly when the caller breaks.
+	count := 0
+	for v := range CountStep(1, 1) {
+		count++
+		if v == 5 {
+			break
+		}
+	}
+	if count != 5 {
+		t.Errorf("expected 5 iterations, got %d", count)
+	}
+}
diff --git a/itertools/go.mod b/itertools/go.mod
new file mode 100644
index 0000000..8b54b0e
--- /dev/null
+++ b/itertools/go.mod
@@ -0,0 +1,3 @@
+module ramalho.org/itertools
+
+go 1.26.2
-- 
2.54.0