`result.partition` test and implementation

2 files changed, 66 insertions, 0 deletions

diff --git a/src/gleam/result.gleam b/src/gleam/result.gleam
index 61e99e9..100d36d 100644
--- a/src/gleam/result.gleam
+++ b/src/gleam/result.gleam
@@ -366,6 +366,36 @@ pub fn all(results: List(Result(a, e))) -> Result(List(a), e) {
   list.try_map(results, fn(x) { x })
 }
 
+/// TODO Add doc!
+pub fn partition(results: List(Result(a, e))) -> Result(List(a), List(e)) {
+  case results {
+    [] -> Ok([])
+    [Ok(a), ..rest] -> do_partition(rest, Ok([a]))
+    [Error(b), ..rest] -> do_partition(rest, Error([b]))
+  }
+}
+
+fn do_partition(results: List(Result(a, b)), acc: Result(List(a), List(b))) {
+  case results {
+    [] ->
+      acc
+      |> map(list.reverse)
+      |> map_error(list.reverse)
+
+    [Ok(a), ..rest] ->
+      case acc {
+        Ok(all_as) -> do_partition(rest, Ok([a, ..all_as]))
+        Error(all_bs) -> do_partition(rest, Error(all_bs))
+      }
+
+    [Error(b), ..rest] ->
+      case acc {
+        Ok(_) -> do_partition(rest, Error([b]))
+        Error(all_bs) -> do_partition(rest, Error([b, ..all_bs]))
+      }
+  }
+}
+
 /// Replace the value within a result
 ///
 /// ## Examples
diff --git a/test/gleam/result_test.gleam b/test/gleam/result_test.gleam
index 6ff35ee..b35034b 100644
--- a/test/gleam/result_test.gleam
+++ b/test/gleam/result_test.gleam
@@ -1,3 +1,4 @@
+import gleam/list
 import gleam/result
 import gleam/should
 
@@ -195,6 +196,41 @@ pub fn all_test() {
   |> should.equal(Error("a"))
 }
 
+pub fn partition_test() {
+  []
+  |> result.partition
+  |> should.equal(Ok([]))
+
+  [Ok(1), Ok(2), Ok(3)]
+  |> result.partition
+  |> should.equal(Ok([1, 2, 3]))
+
+  [Error("a"), Error("b"), Error("c")]
+  |> result.partition
+  |> should.equal(Error(["a", "b", "c"]))
+
+  [Error("a"), Ok(1), Ok(2)]
+  |> result.partition
+  |> should.equal(Error(["a"]))
+
+  [Ok(1), Ok(2), Error("a")]
+  |> result.partition
+  |> should.equal(Error(["a"]))
+
+  [Ok(1), Error("a"), Ok(2), Error("b"), Error("c")]
+  |> result.partition
+  |> should.equal(Error(["a", "b", "c"]))
+
+  // TCO test
+  list.repeat(Ok(1), 1_000_000)
+  |> result.partition
+  |> should.be_ok
+
+  list.repeat(Error("a"), 1_000_000)
+  |> result.partition
+  |> should.be_error
+}
+
 pub fn replace_error_test() {
   Error(Nil)
   |> result.replace_error("Invalid")