[ruby/csv] Fix CSV.filter to preserve headers (#174)

BurdetteLamar · kou · commit 614afb1647d9 · 2020-11-24T09:33:55.000+09:00
Co-authored-by: Sutou Kouhei <kou@clear-code.com> ruby/csv@203c5e0574
diff --git a/lib/csv.rb b/lib/csv.rb
@@ -1057,10 +1057,29 @@ def filter(input=nil, output=nil, **options)
           out_options[key] = value
         end
       end
+
       # build input and output wrappers
-      input  = new(input  || ARGF,    **in_options)
+      input  = new(input  || ARGF, **in_options)
       output = new(output || $stdout, **out_options)
 
+      # process headers
+      need_manual_header_output =
+        (in_options[:headers] and
+         out_options[:headers] == true and
+         out_options[:write_headers])
+      if need_manual_header_output
+        first_row = input.shift
+        if first_row
+          if first_row.is_a?(Row)
+            headers = first_row.headers
+            yield headers
+            output << headers
+          end
+          yield first_row
+          output << first_row
+        end
+      end
+
       # read, yield, write
       input.each do |row|
         yield row
diff --git a/test/csv/interface/test_read_write.rb b/test/csv/interface/test_read_write.rb
@@ -6,7 +6,7 @@ class TestCSVInterfaceReadWrite < Test::Unit::TestCase
   extend DifferentOFS
 
   def test_filter
-    input = <<-CSV
+    input = <<-CSV.freeze
 1;2;3
 4;5
     CSV
@@ -24,6 +24,71 @@ def test_filter
     CSV
   end
 
+  def test_filter_headers_true
+    input = <<-CSV.freeze
+Name,Value
+foo,0
+bar,1
+baz,2
+    CSV
+    output = ""
+    CSV.filter(input, output, headers: true) do |row|
+      row[0] += "X"
+      row[1] = row[1].to_i + 1
+    end
+    assert_equal(<<-CSV, output)
+fooX,1
+barX,2
+bazX,3
+    CSV
+  end
+
+  def test_filter_headers_true_write_headers
+    input = <<-CSV.freeze
+Name,Value
+foo,0
+bar,1
+baz,2
+    CSV
+    output = ""
+    CSV.filter(input, output, headers: true, out_write_headers: true) do |row|
+      if row.is_a?(Array)
+        row[0] += "X"
+        row[1] += "Y"
+      else
+        row[0] += "X"
+        row[1] = row[1].to_i + 1
+      end
+    end
+    assert_equal(<<-CSV, output)
+NameX,ValueY
+fooX,1
+barX,2
+bazX,3
+    CSV
+  end
+
+  def test_filter_headers_array_write_headers
+    input = <<-CSV.freeze
+foo,0
+bar,1
+baz,2
+    CSV
+    output = ""
+    CSV.filter(input, output,
+               headers: ["Name", "Value"],
+               out_write_headers: true) do |row|
+      row[0] += "X"
+      row[1] = row[1].to_i + 1
+    end
+    assert_equal(<<-CSV, output)
+Name,Value
+fooX,1
+barX,2
+bazX,3
+    CSV
+  end
+
   def test_instance_same
     data = ""
     assert_equal(CSV.instance(data, col_sep: ";").object_id,
