Skip to content

Commit

Permalink
[ruby/csv] Fix CSV.filter to preserve headers (#174)
Browse files Browse the repository at this point in the history
Co-authored-by: Sutou Kouhei <kou@clear-code.com>

ruby/csv@203c5e0574
  • Loading branch information
BurdetteLamar authored and kou committed Nov 24, 2020
1 parent 207f2ac commit 614afb1
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 2 deletions.
21 changes: 20 additions & 1 deletion lib/csv.rb
Expand Up @@ -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
Expand Down
67 changes: 66 additions & 1 deletion test/csv/interface/test_read_write.rb
Expand Up @@ -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
Expand All @@ -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,
Expand Down

0 comments on commit 614afb1

Please sign in to comment.