[ruby/stringio] stringio: encoding support

https://github.com/ruby/stringio/commit/7b20075ab0

2 files changed, 31 insertions, 28 deletions

diff --git a/ext/stringio/stringio.c b/ext/stringio/stringio.c
index b02fb113ba..dd5b4c8c4d 100644
--- a/ext/stringio/stringio.c
+++ b/ext/stringio/stringio.c
@@ -190,40 +190,33 @@ strio_initialize(int argc, VALUE *argv, VALUE self)
 static VALUE
 strio_init(int argc, VALUE *argv, struct StringIO *ptr, VALUE self)
 {
-    VALUE string, mode;
-    int trunc = 0;
+    VALUE string, vmode, opt;
+    int oflags;
+    struct rb_io_enc_t convconfig;
 
-    switch (rb_scan_args(argc, argv, "02", &string, &mode)) {
-      case 2:
-	if (FIXNUM_P(mode)) {
-	    int flags = FIX2INT(mode);
-	    ptr->flags = rb_io_oflags_fmode(flags);
-	    trunc = flags & O_TRUNC;
-	}
-	else {
-	    const char *m = StringValueCStr(mode);
-	    ptr->flags = rb_io_modestr_fmode(m);
-	    trunc = *m == 'w';
-	}
+    argc = rb_scan_args(argc, argv, "02:", &string, &vmode, &opt);
+    rb_io_extract_modeenc(&vmode, 0, opt, &oflags, &ptr->flags, &convconfig);
+    if (argc) {
 	StringValue(string);
-	if ((ptr->flags & FMODE_WRITABLE) && OBJ_FROZEN(string)) {
+    }
+    else {
+	string = rb_enc_str_new("", 0, rb_default_external_encoding());
+    }
+    if (OBJ_FROZEN_RAW(string)) {
+	if (ptr->flags & FMODE_WRITABLE) {
 	    rb_syserr_fail(EACCES, 0);
 	}
-	if (trunc) {
-	    rb_str_resize(string, 0);
+    }
+    else {
+	if (NIL_P(vmode)) {
+	    ptr->flags |= FMODE_WRITABLE;
 	}
-	break;
-      case 1:
-	StringValue(string);
-	ptr->flags = OBJ_FROZEN(string) ? FMODE_READABLE : FMODE_READWRITE;
-	break;
-      case 0:
-	string = rb_enc_str_new("", 0, rb_default_external_encoding());
-	ptr->flags = FMODE_READWRITE;
-	break;
+    }
+    if (ptr->flags & FMODE_TRUNC) {
+	rb_str_resize(string, 0);
     }
     ptr->string = string;
-    ptr->enc = 0;
+    ptr->enc = convconfig.enc;
     ptr->pos = 0;
     ptr->lineno = 0;
     RBASIC(self)->flags |= (ptr->flags & FMODE_READWRITE) * (STRIO_READABLE / FMODE_READABLE);
diff --git a/test/stringio/test_stringio.rb b/test/stringio/test_stringio.rb
index 873f67b3b8..1e8e5482ce 100644
--- a/test/stringio/test_stringio.rb
+++ b/test/stringio/test_stringio.rb
@@ -22,7 +22,6 @@ class TestStringIO < Test::Unit::TestCase
     assert_raise(ArgumentError) { StringIO.new('', 'rx') }
     assert_raise(ArgumentError) { StringIO.new('', 'rbt') }
     assert_raise(TypeError) { StringIO.new(nil) }
-    assert_raise(TypeError) { StringIO.new('str', nil) }
 
     o = Object.new
     def o.to_str
@@ -785,6 +784,17 @@ class TestStringIO < Test::Unit::TestCase
     end;
   end
 
+  def test_encoding_write
+    s = StringIO.new("", "w:utf-32be")
+    s.print "abc"
+    assert_equal("abc".encode("utf-32be"), s.string)
+  end
+
+  def test_encoding_read
+    s = StringIO.new("abc".encode("utf-32be"), "r:utf-8")
+    assert_equal("\0\0\0a\0\0\0b\0\0\0c", s.read)
+  end
+
   def assert_string(content, encoding, str, mesg = nil)
     assert_equal([content, encoding], [str, str.encoding], mesg)
   end