ruby · dball · Dec 27, 2012
diff --git a/lib/csv.rb b/lib/csv.rb
@@ -274,7 +274,7 @@ def headers
     #   field( header, offset )
     #   field( index )
     #
-    # This method will fetch the field value by +header+ or +index+.  If a field
+    # This method will return the field value by +header+ or +index+.  If a field
     # is not found, +nil+ is returned.
     #
     # When provided, +offset+ ensures that a header match occurrs on or later
@@ -291,6 +291,39 @@ def field(header_or_index, minimum_index = 0)
     end
     alias_method :[], :field
 
+    #
+    # :call-seq:
+    #   fetch( header )
+    #   fetch( header ) { |row| ... }
+    #   fetch( header, default )
+    #
+    # This method will fetch the field value by +header+. It has the same
+    # behavior as Hash#fetch: if there is a field with the given +header+, its
+    # value is returned. Otherwise, if a block is given, it is yielded the
+    # +header+ and its result is returned; if a +default+ is given as the
+    # second argument, it is returned; otherwise a KeyError is raised.
+    #
+    def fetch(header, *varargs)
+      raise ArgumentError, "Too many arguments" if varargs.length > 1
+      pair = @row.assoc(header)
+      if pair
+        pair.last
+      else
+        if block_given?
+          yield header
+        elsif varargs.empty?
+          raise KeyError, "key not found: #{header}"
+        else
+          varargs.first
+        end
+      end
+    end
+
+    # Returns +true+ if there is a field with the given +header+.
+    def has_key?(header)
+      !!@row.assoc(header)
+    end
+
     #
     # :call-seq:
     #   []=( header, value )

diff --git a/test/csv/test_row.rb b/test/csv/test_row.rb
@@ -78,6 +78,33 @@ def test_field
     assert_equal(nil, @row.field("A", 5))
   end
 
+  def test_fetch
+    # only by name
+    assert_equal(2, @row.fetch('B'))
+
+    # missing header raises KeyError
+    assert_raise KeyError do
+      @row.fetch('foo')
+    end
+
+    # missing header yields itself to block
+    assert_equal 'bar', @row.fetch('foo') { |header|
+      header == 'foo' ? 'bar' : false }
+
+    # missing header returns the given default value
+    assert_equal 'bar', @row.fetch('foo', 'bar')
+
+    # more than one vararg raises ArgumentError
+    assert_raise ArgumentError do
+      @row.fetch('foo', 'bar', 'baz')
+    end
+  end
+
+  def test_has_key?
+    assert_equal(true, @row.has_key?('B'))
+    assert_equal(false, @row.has_key?('foo'))
+  end
+
   def test_set_field
     # set field by name
     assert_equal(100, @row["A"] = 100)