add Enumerable#take, Enumerable#take_while

2 files changed, 51 insertions, 0 deletions

diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb
index bae8f3865..a9545da98 100644
--- a/mrbgems/mruby-enum-ext/mrblib/enum.rb
+++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb
@@ -40,5 +40,46 @@ module Enumerable
     end
     ary
   end
+
+  ##
+  # call-seq:
+  #    enum.take(n)               -> array
+  #
+  # Returns first n elements from <i>enum</i>.
+  #
+  #    a = [1, 2, 3, 4, 5, 0]
+  #    a.take(3)             #=> [1, 2, 3]
+
+  def take(n)
+    raise TypeError, "expected Integer for 1st argument" unless n.kind_of? Integer
+    raise ArgumentError, "attempt to take negative size" if n < 0
+
+    ary = []
+    self.each do |e|
+      break if ary.size >= n
+      ary << e
+    end
+    ary
+  end
+
+  ##
+  # call-seq:
+  #    enum.take_while {|arr| block }   -> array
+  #
+  # Passes elements to the block until the block returns +nil+ or +false+,
+  # then stops iterating and returns an array of all prior elements.
+  #
+  #
+  #    a = [1, 2, 3, 4, 5, 0]
+  #    a.take_while {|i| i < 3 }   #=> [1, 2]
+
+  def take_while(&block)
+    ary = []
+    self.each do |e|
+      return ary unless block.call(e)
+      ary << e
+    end
+    ary
+  end
   
 end
diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb
index 9ab508a11..0f7708db2 100644
--- a/mrbgems/mruby-enum-ext/test/enum.rb
+++ b/mrbgems/mruby-enum-ext/test/enum.rb
@@ -13,3 +13,13 @@ assert("Enumrable#drop_while") do
   assert_equal a.drop_while {|i| i < 3 }, [3, 4, 5, 0]
 end
 
+assert("Enumrable#take") do
+  a = [1, 2, 3, 4, 5, 0]
+  assert_equal a.take(3), [1, 2, 3]
+end
+
+assert("Enumrable#take_while") do
+  a = [1, 2, 3, 4, 5, 0]
+  assert_equal a.take_while {|i| i < 3 }, [1, 2]
+end
+