[ruby/csv] Ensuring StringIO's encoding in CSV.generate (#111)

https://github.com/ruby/csv/commit/dbf55ef008

4 files changed, 38 insertions, 25 deletions

diff --git a/lib/csv.rb b/lib/csv.rb
index 8aa65868b7..c793b8e96a 100644
--- a/lib/csv.rb
+++ b/lib/csv.rb
@@ -531,12 +531,13 @@ class CSV
     # plan to output non-ASCII compatible data.
     #
     def generate(str=nil, **options)
+      encoding = options[:encoding]
       # add a default empty String, if none was given
       if str
         str = StringIO.new(str)
         str.seek(0, IO::SEEK_END)
+        str.set_encoding(encoding) if encoding
       else
-        encoding = options[:encoding]
         str = +""
         str.force_encoding(encoding) if encoding
       end
diff --git a/test/csv/helper.rb b/test/csv/helper.rb
index eadff54408..1f9cf96979 100644
--- a/test/csv/helper.rb
+++ b/test/csv/helper.rb
@@ -15,4 +15,28 @@ module Helper
       ENV["CSV_PARSER_SCANNER_TEST_CHUNK_SIZE"] = chunk_size_keep
     end
   end
+
+  def with_verbose(verbose)
+    original = $VERBOSE
+    begin
+      $VERBOSE = verbose
+      yield
+    ensure
+      $VERBOSE = original
+    end
+  end
+
+  def with_default_internal(encoding)
+    original = Encoding.default_internal
+    begin
+      with_verbose(false) do
+        Encoding.default_internal = encoding
+      end
+      yield
+    ensure
+      with_verbose(false) do
+        Encoding.default_internal = original
+      end
+    end
+  end
 end
diff --git a/test/csv/test_encodings.rb b/test/csv/test_encodings.rb
index acee03db45..cd63af6a83 100644
--- a/test/csv/test_encodings.rb
+++ b/test/csv/test_encodings.rb
@@ -5,6 +5,7 @@ require_relative "helper"
 
 class TestCSVEncodings < Test::Unit::TestCase
   extend DifferentOFS
+  include Helper
 
   def setup
     super
@@ -249,6 +250,15 @@ class TestCSVEncodings < Test::Unit::TestCase
     assert_equal(["foo,\u3042\n".encode(Encoding::Windows_31J), Encoding::Windows_31J], [s, s.encoding], bug9766)
   end
 
+  def test_encoding_with_default_internal
+    with_default_internal(Encoding::UTF_8) do
+      s = CSV.generate(String.new(encoding: Encoding::Big5), encoding: Encoding::Big5) do |csv|
+        csv << ["漢字"]
+      end
+      assert_equal(["漢字\n".encode(Encoding::Big5), Encoding::Big5], [s, s.encoding])
+    end
+  end
+
   def test_row_separator_detection_with_invalid_encoding
     csv = CSV.new("invalid,\xF8\r\nvalid,x\r\n".force_encoding("UTF-8"),
                   encoding: "UTF-8")
diff --git a/test/csv/write/test_general.rb b/test/csv/write/test_general.rb
index bdc8aadd55..677119e1ae 100644
--- a/test/csv/write/test_general.rb
+++ b/test/csv/write/test_general.rb
@@ -4,6 +4,8 @@
 require_relative "../helper"
 
 module TestCSVWriteGeneral
+  include Helper
+
   def test_tab
     assert_equal("\t#{$INPUT_RECORD_SEPARATOR}",
                  generate_line(["\t"]))
@@ -221,30 +223,6 @@ module TestCSVWriteGeneral
                    generate_line(row))
     end
   end
-
-  def with_verbose(verbose)
-    original = $VERBOSE
-    begin
-      $VERBOSE = verbose
-      yield
-    ensure
-      $VERBOSE = original
-    end
-  end
-
-  def with_default_internal(encoding)
-    original = Encoding.default_internal
-    begin
-      with_verbose(false) do
-        Encoding.default_internal = encoding
-      end
-      yield
-    ensure
-      with_verbose(false) do
-        Encoding.default_internal = original
-      end
-    end
-  end
 end
 
 class TestCSVWriteGeneralGenerateLine < Test::Unit::TestCase