diff options
Diffstat (limited to 'mrbgems/mruby-enum-ext')
| -rw-r--r-- | mrbgems/mruby-enum-ext/mrblib/enum.rb | 41 | ||||
| -rw-r--r-- | mrbgems/mruby-enum-ext/test/enum.rb | 10 |
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 + |
