From 5d1608fae748069860f315eb42aa7bc92a5e763b Mon Sep 17 00:00:00 2001
From: Philip Ross <dev@baseballlover723.com>
Date: Thu, 15 Jun 2023 13:57:45 -0700
Subject: [PATCH] added a partition overload to accept a class

---
 spec/std/enumerable_spec.cr | 32 ++++++++++++++++++++++++++++++++
 src/enumerable.cr           | 23 +++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/spec/std/enumerable_spec.cr b/spec/std/enumerable_spec.cr
index 435a056191e3..3b3cef2bd0c1 100644
--- a/spec/std/enumerable_spec.cr
+++ b/spec/std/enumerable_spec.cr
@@ -1049,6 +1049,38 @@ describe "Enumerable" do
   describe "partition" do
     it { [1, 2, 2, 3].partition { |x| x == 2 }.should eq({[2, 2], [1, 3]}) }
     it { [1, 2, 3, 4, 5, 6].partition(&.even?).should eq({[2, 4, 6], [1, 3, 5]}) }
+
+    it "with mono type on union type" do
+      ints, others = [1, true, nil, 3, false, "string", 'c'].partition(Int32)
+      ints.should eq([1, 3])
+      others.should eq([true, nil, false, "string", 'c'])
+      ints.should be_a(Array(Int32))
+      others.should be_a(Array(Bool | String | Char | Nil))
+    end
+
+    it "with union type on union type" do
+      ints_bools, others = [1, true, nil, 3, false, "string", 'c'].partition(Int32 | Bool)
+      ints_bools.should eq([1, true, 3, false])
+      others.should eq([nil, "string", 'c'])
+      ints_bools.should be_a(Array(Int32 | Bool))
+      others.should be_a(Array(String | Char | Nil))
+    end
+
+    it "with missing type on union type" do
+      symbols, others = [1, true, nil, 3, false, "string", 'c'].partition(Symbol)
+      symbols.empty?.should be_true
+      others.should eq([1, true, nil, 3, false, "string", 'c'])
+      symbols.should be_a(Array(Symbol))
+      others.should be_a(Array(Int32 | Bool | String | Char | Nil))
+    end
+
+    it "with mono type on mono type" do
+      ints, others = [1, 3].partition(Int32)
+      ints.should eq([1, 3])
+      others.empty?.should be_true
+      ints.should be_a(Array(Int32))
+      others.should be_a(Array(NoReturn))
+    end
   end
 
   describe "reject" do
diff --git a/src/enumerable.cr b/src/enumerable.cr
index d7f978bd088d..7d95cb6f81b2 100644
--- a/src/enumerable.cr
+++ b/src/enumerable.cr
@@ -1441,6 +1441,29 @@ module Enumerable(T)
     {a, b}
   end
 
+  # Returns a `Tuple` with two arrays. The first one contains the elements
+  # in the collection that are of the given *type*,
+  # and the second one that are **not** of the given *type*.
+  #
+  # ```
+  # ints, others = [1, true, nil, 3, false, "string", 'c'].partition(Int32)
+  # ints           # => [1, 3]
+  # others         # => [true, nil, false, "string", 'c']
+  # typeof(ints)   # => Array(Int32)
+  # typeof(others) # => Array(String | Char | Nil)
+  # ```
+  def partition(type : U.class) forall U
+    a = [] of U
+    b = [] of typeof(begin
+      e = Enumerable.element_type(self)
+      e.is_a?(U) ? raise("") : e
+    end)
+    each do |e|
+      e.is_a?(U) ? a.push(e) : b.push(e)
+    end
+    {a, b}
+  end
+
   # Returns an `Array` with all the elements in the collection for which
   # the passed block is falsey.
   #