diff options
| -rw-r--r-- | lib/csv.rb | 21 | ||||
| -rw-r--r-- | 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, |