diff options
| author | Burdette Lamar <BurdetteLamar@Yahoo.com> | 2020-09-11 16:36:01 -0500 |
|---|---|---|
| committer | Sutou Kouhei <kou@cozmixng.org> | 2020-11-24 09:33:55 +0900 |
| commit | 614afb1647d9c9eb170262c8b033f000c5beb6f0 (patch) | |
| tree | d8edc7ead9f0bb31ba8ccb806b83e6feef9216ac | |
| parent | 207f2acc1355dea1fc1f483e4d8ff3e571a0ad89 (diff) | |
| download | ruby-614afb1647d9c9eb170262c8b033f000c5beb6f0.tar.gz | |
[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
| -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, |