Add Array#fetch

2 files changed, 52 insertions, 0 deletions

diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index feec10ead..257862b2b 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -211,4 +211,45 @@ class Array
     end
     self
   end
+
+  NONE=Object.new
+  ##
+  #  call-seq:
+  #     ary.fetch(index)                    -> obj
+  #     ary.fetch(index, default)           -> obj
+  #     ary.fetch(index) { |index| block }  -> obj
+  #
+  #  Tries to return the element at position +index+, but throws an IndexError
+  #  exception if the referenced +index+ lies outside of the array bounds.  This
+  #  error can be prevented by supplying a second argument, which will act as a
+  #  +default+ value.
+  #
+  #  Alternatively, if a block is given it will only be executed when an
+  #  invalid +index+ is referenced.  Negative values of +index+ count from the
+  #  end of the array.
+  #
+  #     a = [ 11, 22, 33, 44 ]
+  #     a.fetch(1)               #=> 22
+  #     a.fetch(-1)              #=> 44
+  #     a.fetch(4, 'cat')        #=> "cat"
+  #     a.fetch(100) { |i| puts "#{i} is out of bounds" }
+  #                              #=> "100 is out of bounds"
+  #
+
+  def fetch(n=nil, ifnone=NONE, &block)
+    warn "block supersedes default value argument" if n != nil && ifnone != NONE && block
+
+    idx = n
+    if idx < 0
+      idx += size
+    end
+    if idx < 0 || size <= idx
+      return block.call(n) if block
+      if ifnone == NONE
+        raise IndexError, "index #{n} outside of array bounds: #{-size}...#{size}"
+      end
+      return ifnone
+    end
+    self[idx]
+  end
 end
diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb
index 1c441cec4..ab830cca7 100644
--- a/mrbgems/mruby-array-ext/test/array.rb
+++ b/mrbgems/mruby-array-ext/test/array.rb
@@ -107,3 +107,14 @@ assert("Array#compact!") do
   a.compact!
   a == [1, "2", :t, false]
 end
+
+assert("Array#fetch") do
+  a = [ 11, 22, 33, 44 ]
+  assert_equal 22, a.fetch(1)
+  assert_equal 44, a.fetch(-1)
+  assert_equal 'cat', a.fetch(4, 'cat')
+  ret = 0
+  a.fetch(100) { |i| ret = i }
+  assert_equal 100, ret
+  assert_raise(IndexError) { a.fetch(100) }
+end