force_quotes: add support for specifying the target indexes or names

kou · kou · commit 8812c58a2643 · 2020-07-16T06:10:38.000+09:00
GitHub: fix GH-153 Reported by Aleksandr. Thanks!!!
diff --git a/lib/csv/writer.rb b/lib/csv/writer.rb
@@ -43,8 +43,10 @@ def <<(row)
 
       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 @@ def prepare_header
       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 @@ def prepare_format
         @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 @@ def quote_field(field)
         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
@@ -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
