diff options
| -rw-r--r-- | lib/csv/writer.rb | 47 | ||||
| -rw-r--r-- | test/csv/write/test_force_quotes.rb | 78 |
2 files changed, 122 insertions, 3 deletions
diff --git a/lib/csv/writer.rb b/lib/csv/writer.rb index 6c02d7473e..d49115fccf 100644 --- a/lib/csv/writer.rb +++ b/lib/csv/writer.rb @@ -43,8 +43,10 @@ class CSV row = @fields_converter.convert(row, nil, lineno) if @fields_converter + i = -1 converted_row = row.collect do |field| - quote(field) + i += 1 + quote(field, i) end line = converted_row.join(@column_separator) + @row_separator if @output_encoding @@ -100,6 +102,33 @@ class CSV end end + def prepare_force_quotes_fields(force_quotes) + @force_quotes_fields = {} + force_quotes.each do |name_or_index| + case name_or_index + when Integer + index = name_or_index + @force_quotes_fields[index] = true + when String, Symbol + name = name_or_index.to_s + if @headers.nil? + message = ":headers is required when you use field name " + + "in :force_quotes: " + + "#{name_or_index.inspect}: #{force_quotes.inspect}" + raise ArgumentError, message + end + index = @headers.index(name) + next if index.nil? + @force_quotes_fields[index] = true + else + message = ":force_quotes element must be " + + "field index or field name: " + + "#{name_or_index.inspect}: #{force_quotes.inspect}" + raise ArgumentError, message + end + end + end + def prepare_format @column_separator = @options[:column_separator].to_s.encode(@encoding) row_separator = @options[:row_separator] @@ -109,7 +138,17 @@ class CSV @row_separator = row_separator.to_s.encode(@encoding) end @quote_character = @options[:quote_character] - @force_quotes = @options[:force_quotes] + force_quotes = @options[:force_quotes] + if force_quotes.is_a?(Array) + prepare_force_quotes_fields(force_quotes) + @force_quotes = false + elsif force_quotes + @force_quotes_fields = nil + @force_quotes = true + else + @force_quotes_fields = nil + @force_quotes = false + end unless @force_quotes @quotable_pattern = Regexp.new("[\r\n".encode(@encoding) + @@ -147,9 +186,11 @@ class CSV encoded_quote_character end - def quote(field) + def quote(field, i) if @force_quotes quote_field(field) + elsif @force_quotes_fields and @force_quotes_fields[i] + quote_field(field) else if field.nil? # represent +nil+ fields as empty unquoted fields "" diff --git a/test/csv/write/test_force_quotes.rb b/test/csv/write/test_force_quotes.rb new file mode 100644 index 0000000000..622dcb021b --- /dev/null +++ b/test/csv/write/test_force_quotes.rb @@ -0,0 +1,78 @@ +# frozen_string_literal: false + +require_relative "../helper" + +module TestCSVWriteForceQuotes + def test_default + assert_equal(%Q[1,2,3#{$INPUT_RECORD_SEPARATOR}], + generate_line(["1", "2", "3"])) + end + + def test_true + assert_equal(%Q["1","2","3"#{$INPUT_RECORD_SEPARATOR}], + generate_line(["1", "2", "3"], + force_quotes: true)) + end + + def test_false + assert_equal(%Q[1,2,3#{$INPUT_RECORD_SEPARATOR}], + generate_line(["1", "2", "3"], + force_quotes: false)) + end + + def test_field_name + assert_equal(%Q["1",2,"3"#{$INPUT_RECORD_SEPARATOR}], + generate_line(["1", "2", "3"], + headers: ["a", "b", "c"], + force_quotes: ["a", :c])) + end + + def test_field_name_without_headers + force_quotes = ["a", "c"] + error = assert_raise(ArgumentError) do + generate_line(["1", "2", "3"], + force_quotes: force_quotes) + end + assert_equal(":headers is required when you use field name " + + "in :force_quotes: " + + "#{force_quotes.first.inspect}: #{force_quotes.inspect}", + error.message) + end + + def test_field_index + assert_equal(%Q["1",2,"3"#{$INPUT_RECORD_SEPARATOR}], + generate_line(["1", "2", "3"], + force_quotes: [0, 2])) + end + + def test_field_unknown + force_quotes = [1.1] + error = assert_raise(ArgumentError) do + generate_line(["1", "2", "3"], + force_quotes: force_quotes) + end + assert_equal(":force_quotes element must be field index or field name: " + + "#{force_quotes.first.inspect}: #{force_quotes.inspect}", + error.message) + end +end + +class TestCSVWriteForceQuotesGenerateLine < Test::Unit::TestCase + include TestCSVWriteForceQuotes + extend DifferentOFS + + def generate_line(row, **kwargs) + CSV.generate_line(row, **kwargs) + end +end + +class TestCSVWriteForceQuotesGenerate < Test::Unit::TestCase + include TestCSVWriteForceQuotes + extend DifferentOFS + + def generate_line(row, **kwargs) + CSV.generate(**kwargs) do |csv| + csv << row + end + end +end |