summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorBurdette Lamar <BurdetteLamar@Yahoo.com>2020-09-11 16:36:01 -0500
committerSutou Kouhei <kou@cozmixng.org>2020-11-24 09:33:55 +0900
commit614afb1647d9c9eb170262c8b033f000c5beb6f0 (patch)
treed8edc7ead9f0bb31ba8ccb806b83e6feef9216ac
parent207f2acc1355dea1fc1f483e4d8ff3e571a0ad89 (diff)
[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
Notes
Notes: Merged: https://github.com/ruby/ruby/pull/3804
-rw-r--r--lib/csv.rb21
-rw-r--r--test/csv/interface/test_read_write.rb67
2 files changed, 86 insertions, 2 deletions
diff --git a/lib/csv.rb b/lib/csv.rb
index 143909a..cfd1b86 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 877e5f3..20c9fe3 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,