From 614afb1647d9c9eb170262c8b033f000c5beb6f0 Mon Sep 17 00:00:00 2001
From: Burdette Lamar <BurdetteLamar@Yahoo.com>
Date: Fri, 11 Sep 2020 16:36:01 -0500
Subject: [ruby/csv] Fix CSV.filter to preserve headers (#174)

Co-authored-by: Sutou Kouhei <kou@clear-code.com>

https://github.com/ruby/csv/commit/203c5e0574
---
 lib/csv.rb                            | 21 ++++++++++-
 test/csv/interface/test_read_write.rb | 67 ++++++++++++++++++++++++++++++++++-
 2 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/lib/csv.rb b/lib/csv.rb
index 143909a099..cfd1b8621e 100644
--- a/lib/csv.rb
+++ b/lib/csv.rb
@@ -1057,10 +1057,29 @@ class CSV
           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
index 877e5f355e..20c9fe317e 100644
--- 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 @@ class TestCSVInterfaceReadWrite < Test::Unit::TestCase
     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,
-- 
cgit v1.2.3