Merge pull request #33137 from bogdanvlviv/add-array-extract-method

Add `Array#extract!`

9 files changed, 107 insertions, 11 deletions

diff --git a/actionpack/lib/action_dispatch/http/parameter_filter.rb b/actionpack/lib/action_dispatch/http/parameter_filter.rb
index 1d58964862..de11939fa8 100644
--- a/actionpack/lib/action_dispatch/http/parameter_filter.rb
+++ b/actionpack/lib/action_dispatch/http/parameter_filter.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "active_support/core_ext/object/duplicable"
+require "active_support/core_ext/array/extract"
 
 module ActionDispatch
   module Http
@@ -38,8 +39,8 @@ module ActionDispatch
             end
           end
 
-          deep_regexps, regexps = regexps.partition { |r| r.to_s.include?("\\.".freeze) }
-          deep_strings, strings = strings.partition { |s| s.include?("\\.".freeze) }
+          deep_regexps = regexps.extract! { |r| r.to_s.include?("\\.".freeze) }
+          deep_strings = strings.extract! { |s| s.include?("\\.".freeze) }
 
           regexps << Regexp.new(strings.join("|".freeze), true) unless strings.empty?
           deep_regexps << Regexp.new(deep_strings.join("|".freeze), true) unless deep_strings.empty?
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/oid/type_map_initializer.rb b/activerecord/lib/active_record/connection_adapters/postgresql/oid/type_map_initializer.rb
index 231278c184..79351bc3a4 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/oid/type_map_initializer.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/oid/type_map_initializer.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/array/extract"
+
 module ActiveRecord
   module ConnectionAdapters
     module PostgreSQL
@@ -16,12 +18,12 @@ module ActiveRecord
 
           def run(records)
             nodes = records.reject { |row| @store.key? row["oid"].to_i }
-            mapped, nodes = nodes.partition { |row| @store.key? row["typname"] }
-            ranges, nodes = nodes.partition { |row| row["typtype"] == "r".freeze }
-            enums, nodes = nodes.partition { |row| row["typtype"] == "e".freeze }
-            domains, nodes = nodes.partition { |row| row["typtype"] == "d".freeze }
-            arrays, nodes = nodes.partition { |row| row["typinput"] == "array_in".freeze }
-            composites, nodes = nodes.partition { |row| row["typelem"].to_i != 0 }
+            mapped = nodes.extract! { |row| @store.key? row["typname"] }
+            ranges = nodes.extract! { |row| row["typtype"] == "r".freeze }
+            enums = nodes.extract! { |row| row["typtype"] == "e".freeze }
+            domains = nodes.extract! { |row| row["typtype"] == "d".freeze }
+            arrays = nodes.extract! { |row| row["typinput"] == "array_in".freeze }
+            composites = nodes.extract! { |row| row["typelem"].to_i != 0 }
 
             mapped.each     { |row| register_mapped_type(row)    }
             enums.each      { |row| register_enum_type(row)      }
diff --git a/activerecord/lib/active_record/relation/predicate_builder/array_handler.rb b/activerecord/lib/active_record/relation/predicate_builder/array_handler.rb
index 64bf83e3c1..e5191fa38a 100644
--- a/activerecord/lib/active_record/relation/predicate_builder/array_handler.rb
+++ b/activerecord/lib/active_record/relation/predicate_builder/array_handler.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/array/extract"
+
 module ActiveRecord
   class PredicateBuilder
     class ArrayHandler # :nodoc:
@@ -11,8 +13,8 @@ module ActiveRecord
         return attribute.in([]) if value.empty?
 
         values = value.map { |x| x.is_a?(Base) ? x.id : x }
-        nils, values = values.partition(&:nil?)
-        ranges, values = values.partition { |v| v.is_a?(Range) }
+        nils = values.extract!(&:nil?)
+        ranges = values.extract! { |v| v.is_a?(Range) }
 
         values_predicate =
           case values.length
diff --git a/activesupport/CHANGELOG.md b/activesupport/CHANGELOG.md
index 25e2ee04f9..4ae02edd6a 100644
--- a/activesupport/CHANGELOG.md
+++ b/activesupport/CHANGELOG.md
@@ -1,3 +1,14 @@
+*   Add `Array#extract!`.
+
+    The method removes and returns the elements for which the block returns a true value.
+    If no block is given, an Enumerator is returned instead.
+
+        numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+        odd_numbers = numbers.extract! { |number| number.odd? } # => [1, 3, 5, 7, 9]
+        numbers # => [0, 2, 4, 6, 8]
+
+    *bogdanvlviv*
+
 *   Support not to cache `nil` for `ActiveSupport::Cache#fetch`.
 
         cache.fetch('bar', skip_nil: true) { nil }
diff --git a/activesupport/lib/active_support/core_ext/array.rb b/activesupport/lib/active_support/core_ext/array.rb
index 6d83b76882..a2569c798b 100644
--- a/activesupport/lib/active_support/core_ext/array.rb
+++ b/activesupport/lib/active_support/core_ext/array.rb
@@ -3,6 +3,7 @@
 require "active_support/core_ext/array/wrap"
 require "active_support/core_ext/array/access"
 require "active_support/core_ext/array/conversions"
+require "active_support/core_ext/array/extract"
 require "active_support/core_ext/array/extract_options"
 require "active_support/core_ext/array/grouping"
 require "active_support/core_ext/array/prepend_and_append"
diff --git a/activesupport/lib/active_support/core_ext/array/extract.rb b/activesupport/lib/active_support/core_ext/array/extract.rb
new file mode 100644
index 0000000000..cc5a8a3f88
--- /dev/null
+++ b/activesupport/lib/active_support/core_ext/array/extract.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+class Array
+  # Removes and returns the elements for which the block returns a true value.
+  # If no block is given, an Enumerator is returned instead.
+  #
+  #   numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+  #   odd_numbers = numbers.extract! { |number| number.odd? } # => [1, 3, 5, 7, 9]
+  #   numbers # => [0, 2, 4, 6, 8]
+  def extract!
+    return to_enum(:extract!) { size } unless block_given?
+
+    extracted_elements = []
+
+    reject! do |element|
+      extracted_elements << element if yield(element)
+    end
+
+    extracted_elements
+  end
+end
diff --git a/activesupport/test/core_ext/array/extract_test.rb b/activesupport/test/core_ext/array/extract_test.rb
new file mode 100644
index 0000000000..200727667c
--- /dev/null
+++ b/activesupport/test/core_ext/array/extract_test.rb
@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require "abstract_unit"
+require "active_support/core_ext/array"
+
+class ExtractTest < ActiveSupport::TestCase
+  def test_extract
+    numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    array_id = numbers.object_id
+
+    odd_numbers = numbers.extract!(&:odd?)
+
+    assert_equal [1, 3, 5, 7, 9], odd_numbers
+    assert_equal [0, 2, 4, 6, 8], numbers
+    assert_equal array_id, numbers.object_id
+  end
+
+  def test_extract_without_block
+    numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    array_id = numbers.object_id
+
+    extract_enumerator = numbers.extract!
+
+    assert_instance_of Enumerator, extract_enumerator
+    assert_equal numbers.size, extract_enumerator.size
+
+    odd_numbers = extract_enumerator.each(&:odd?)
+
+    assert_equal [1, 3, 5, 7, 9], odd_numbers
+    assert_equal [0, 2, 4, 6, 8], numbers
+    assert_equal array_id, numbers.object_id
+  end
+
+  def test_extract_on_empty_array
+    empty_array = []
+    array_id = empty_array.object_id
+
+    new_empty_array = empty_array.extract! {}
+
+    assert_equal [], new_empty_array
+    assert_equal [], empty_array
+    assert_equal array_id, empty_array.object_id
+  end
+end
diff --git a/guides/source/active_support_core_extensions.md b/guides/source/active_support_core_extensions.md
index dfd21915b0..f9fc7044ba 100644
--- a/guides/source/active_support_core_extensions.md
+++ b/guides/source/active_support_core_extensions.md
@@ -2156,6 +2156,19 @@ This method is an alias of `Array#<<`.
 
 NOTE: Defined in `active_support/core_ext/array/prepend_and_append.rb`.
 
+### Extracting
+
+The method `extract!` removes and returns the elements for which the block returns a true value.
+If no block is given, an Enumerator is returned instead.
+
+```ruby
+numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+odd_numbers = numbers.extract! { |number| number.odd? } # => [1, 3, 5, 7, 9]
+numbers # => [0, 2, 4, 6, 8]
+```
+
+NOTE: Defined in `active_support/core_ext/array/extract.rb`.
+
 ### Options Extraction
 
 When the last argument in a method call is a hash, except perhaps for a `&block` argument, Ruby allows you to omit the brackets:
diff --git a/railties/lib/rails/api/generator.rb b/railties/lib/rails/api/generator.rb
index 3405560b74..126d4d0438 100644
--- a/railties/lib/rails/api/generator.rb
+++ b/railties/lib/rails/api/generator.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "sdoc"
+require "active_support/core_ext/array/extract"
 
 class RDoc::Generator::API < RDoc::Generator::SDoc # :nodoc:
   RDoc::RDoc.add_generator self
@@ -11,7 +12,7 @@ class RDoc::Generator::API < RDoc::Generator::SDoc # :nodoc:
     # since they aren't nested under a definition of the `ActiveStorage` module.
     if visited.empty?
       classes = classes.reject { |klass| active_storage?(klass) }
-      core_exts, classes = classes.partition { |klass| core_extension?(klass) }
+      core_exts = classes.extract! { |klass| core_extension?(klass) }
 
       super.unshift([ "Core extensions", "", "", build_core_ext_subtree(core_exts, visited) ])
     else