Merge pull request #19 from chef/lcg/stringio-encoding

support encoding StringIOs

3 files changed, 31 insertions, 15 deletions

diff --git a/ext/ffi_yajl/ext/encoder/encoder.c b/ext/ffi_yajl/ext/encoder/encoder.c
index b43d02b..fdc6f56 100644
--- a/ext/ffi_yajl/ext/encoder/encoder.c
+++ b/ext/ffi_yajl/ext/encoder/encoder.c
@@ -2,7 +2,7 @@
 #include <yajl/yajl_gen.h>
 
 static VALUE mFFI_Yajl, mExt, mEncoder, mEncoder2, cEncodeError;
-static VALUE cDate, cTime, cDateTime;
+static VALUE cDate, cTime, cDateTime, cStringIO;
 static VALUE cYajl_Gen;
 
 /* FIXME: the json gem does a whole bunch of indirection around monkeypatching...  not sure if we need to as well... */
@@ -295,11 +295,8 @@ static VALUE rb_cString_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
   return Qnil;
 }
 
-/* calls #to_s on an object to encode it */
-static VALUE object_to_s_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
+static VALUE string_ffi_yajl(VALUE str, VALUE rb_yajl_gen, VALUE state) {
   yajl_gen_status status;
-  ID sym_to_s = rb_intern("to_s");
-  VALUE str = rb_funcall(self, sym_to_s, 0);
   char *cptr = RSTRING_PTR(str);
   int len = RSTRING_LEN(str);
   struct yajl_gen_t *yajl_gen;
@@ -312,6 +309,13 @@ static VALUE object_to_s_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
   return Qnil;
 }
 
+/* calls #to_s on an object to encode it */
+static VALUE object_to_s_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
+  ID sym_to_s = rb_intern("to_s");
+  VALUE str = rb_funcall(self, sym_to_s, 0);
+  return string_ffi_yajl(str, rb_yajl_gen, state);
+}
+
 static VALUE rb_cSymbol_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
   return object_to_s_ffi_yajl(self, rb_yajl_gen, state);
 }
@@ -321,19 +325,15 @@ static VALUE rb_cDate_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
 }
 
 static VALUE rb_cTime_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
-  yajl_gen_status status;
   ID sym_strftime = rb_intern("strftime");
   VALUE str = rb_funcall(self, sym_strftime, 1, rb_str_new2("%Y-%m-%d %H:%M:%S %z"));
-  char *cptr = RSTRING_PTR(str);
-  int len = RSTRING_LEN(str);
-  struct yajl_gen_t *yajl_gen;
-  Data_Get_Struct(rb_yajl_gen, struct yajl_gen_t, yajl_gen);
-
-  CHECK_STATUS(
-    yajl_gen_string(yajl_gen, (unsigned char *)cptr, len)
-  );
+  return string_ffi_yajl(str, rb_yajl_gen, state);
+}
 
-  return Qnil;
+static VALUE rb_cStringIO_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
+  ID sym_read = rb_intern("read");
+  VALUE str = rb_funcall(self, sym_read, 0);
+  return string_ffi_yajl(str, rb_yajl_gen, state);
 }
 
 static VALUE rb_cDateTime_ffi_yajl(VALUE self, VALUE rb_yajl_gen, VALUE state) {
@@ -373,6 +373,7 @@ void Init_encoder() {
   cDate = rb_define_class("Date", rb_cObject);
   cTime = rb_define_class("Time", rb_cObject);
   cDateTime = rb_define_class("DateTime", cDate);
+  cStringIO = rb_define_class("StringIO", rb_cData);
 
   rb_define_method(rb_cHash, "ffi_yajl", rb_cHash_ffi_yajl, 2);
   rb_define_method(rb_cArray, "ffi_yajl", rb_cArray_ffi_yajl, 2);
@@ -387,5 +388,6 @@ void Init_encoder() {
   rb_define_method(cDate, "ffi_yajl", rb_cDate_ffi_yajl, 2);
   rb_define_method(cTime, "ffi_yajl", rb_cTime_ffi_yajl, 2);
   rb_define_method(cDateTime, "ffi_yajl", rb_cDateTime_ffi_yajl, 2);
+  rb_define_method(cStringIO, "ffi_yajl", rb_cStringIO_ffi_yajl, 2);
   rb_define_method(rb_cObject, "ffi_yajl", rb_cObject_ffi_yajl, 2);
 }
diff --git a/lib/ffi_yajl/ffi/encoder.rb b/lib/ffi_yajl/ffi/encoder.rb
index 6e168fe..73a0d11 100644
--- a/lib/ffi_yajl/ffi/encoder.rb
+++ b/lib/ffi_yajl/ffi/encoder.rb
@@ -227,6 +227,15 @@ class String
   end
 end
 
+class StringIO
+  def ffi_yajl(yajl_gen, state)
+    str = self.read
+    if ( status = FFI_Yajl.yajl_gen_string(yajl_gen, str, str.bytesize) ) != 0
+      FFI_Yajl::Encoder.raise_error_for_status(status)
+    end
+  end
+end
+
 class Date
   def ffi_yajl(yajl_gen, state)
     str = self.to_s
diff --git a/spec/ffi_yajl/encoder_spec.rb b/spec/ffi_yajl/encoder_spec.rb
index 2f08e3f..0eef23d 100644
--- a/spec/ffi_yajl/encoder_spec.rb
+++ b/spec/ffi_yajl/encoder_spec.rb
@@ -131,6 +131,11 @@ describe "FFI_Yajl::Encoder" do
     expect(encoder.encode(ruby)).to eq( %q{"2001-02-03"} )
   end
 
+  it "can encode StringIOs" do
+    ruby = { "foo" => StringIO.new('THING') }
+    expect(encoder.encode(ruby)).to eq("{\"foo\":\"THING\"}")
+  end
+
   context "when encoding Time objects in UTC timezone" do
     before do
       @saved_tz = ENV['TZ']