Merge pull request #24 from technorama/master

add symbolize_keys support to unpack (compatible with msgpack-jruby)

Author: frsyuki

ext/msgpack/buffer.h

@@ -108,6 +108,7 @@ struct msgpack_buffer_t {
     size_t write_reference_threshold;
     size_t read_reference_threshold;
     size_t io_buffer_size;
+    bool symbolize_keys;
 
     VALUE owner;
 };
@@ -156,6 +157,11 @@ static inline void msgpack_buffer_reset_io(msgpack_buffer_t* b)
     b->io = Qnil;
 }
 
+static inline void msgpack_buffer_set_symbolize_keys(msgpack_buffer_t* b, bool val)
+{
+    b->symbolize_keys = val;
+}
+
 static inline bool msgpack_buffer_has_io(msgpack_buffer_t* b)
 {
     return b->io != Qnil;
@@ -418,8 +424,14 @@ static inline VALUE _msgpack_buffer_refer_head_mapped_string(msgpack_buffer_t* b
     return rb_str_substr(b->head->mapped_string, offset, length);
 }
 
-static inline VALUE msgpack_buffer_read_top_as_string(msgpack_buffer_t* b, size_t length, bool will_be_frozen)
+static inline VALUE msgpack_buffer_read_top_as_string(msgpack_buffer_t* b, size_t length, bool will_be_frozen, bool symbolize_keys)
 {
+    if(symbolize_keys) {
+        VALUE result = ID2SYM(rb_intern2(b->read_buffer, length));
+        _msgpack_buffer_consumed(b, length);
+        return result;
+    }
+
 #ifndef DISABLE_BUFFER_READ_REFERENCE_OPTIMIZE
     /* optimize */
     if(!will_be_frozen &&

ext/msgpack/buffer_class.c

@@ -105,6 +105,11 @@ void MessagePack_Buffer_initialize(msgpack_buffer_t* b, VALUE io, VALUE options)
         if(v != Qnil) {
             msgpack_buffer_set_io_buffer_size(b, NUM2ULONG(v));
         }
+
+        v = rb_hash_aref(options, ID2SYM(rb_intern("symbolize_keys")));
+        if(v != Qnil) {
+            msgpack_buffer_set_symbolize_keys(b, RTEST(v));
+        }
     }
 }
 

ext/msgpack/unpacker.c

@@ -257,10 +257,15 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, bool str)
         /* don't use zerocopy for hash keys but get a frozen string directly
          * because rb_hash_aset freezes keys and it causes copying */
         bool will_freeze = is_reading_map_key(uk);
-        VALUE string = msgpack_buffer_read_top_as_string(UNPACKER_BUFFER_(uk), length, will_freeze);
-        object_complete_string(uk, string);
-        if(will_freeze) {
-            rb_obj_freeze(string);
+        bool symbolize_keys = will_freeze && UNPACKER_BUFFER_(uk)->symbolize_keys;
+        VALUE string = msgpack_buffer_read_top_as_string(UNPACKER_BUFFER_(uk), length, will_freeze, symbolize_keys);
+        if(symbolize_keys) {
+            object_complete(uk, string);
+        } else {
+            object_complete_string(uk, string);
+            if(will_freeze) {
+                rb_obj_freeze(string);
+            }
         }
         uk->reading_raw_remaining = 0;
         return PRIMITIVE_OBJECT_COMPLETE;

ext/msgpack/unpacker_class.c

@@ -288,11 +288,16 @@ static VALUE Unpacker_reset(VALUE self)
 VALUE MessagePack_unpack(int argc, VALUE* argv)
 {
     VALUE src;
+    VALUE options = Qnil;
 
     switch(argc) {
     case 1:
         src = argv[0];
         break;
+    case 2:
+        src = argv[0];
+        options = argv[1];
+        break;
     default:
         rb_raise(rb_eArgError, "wrong number of arguments (%d for 1)", argc);
     }
@@ -312,11 +317,12 @@ VALUE MessagePack_unpack(int argc, VALUE* argv)
     msgpack_buffer_set_write_reference_threshold(UNPACKER_BUFFER_(uk), 0);
 
     if(io != Qnil) {
-        MessagePack_Buffer_initialize(UNPACKER_BUFFER_(uk), io, Qnil);
+        MessagePack_Buffer_initialize(UNPACKER_BUFFER_(uk), io, options);
     }
 
     if(src != Qnil) {
         /* prefer reference than copying; see MessagePack_Unpacker_module_init */
+        MessagePack_Buffer_initialize(UNPACKER_BUFFER_(uk), io, options);
         msgpack_buffer_append_string(UNPACKER_BUFFER_(uk), src);
     }
 

spec/buffer_io_spec.rb

@@ -201,7 +201,7 @@ def set_source(s)
       fragments.each {|s|
         lambda {
           b.skip_all(s.size)
-        }.should_not raise_error(EOFError)
+        }.should_not raise_error
       }
       b.empty?.should == true
       lambda {

spec/unpacker_spec.rb

@@ -230,6 +230,14 @@
     parsed.should == true
   end
 
+  it 'symbolize_keys' do
+    hash = { a: :b, c: :d }
+    hash_result = { a: 'b', c: 'd' }
+    MessagePack.unpack(MessagePack.pack(hash), symbolize_keys: true).should == hash_result
+    hash = { 'a' => 'b', 'c' => 'd' }
+    MessagePack.unpack(MessagePack.pack(hash), symbolize_keys: true).should == hash_result
+  end
+
   it "msgpack str 8 type" do
     MessagePack.unpack([0xd9, 0x00].pack('C*')).should == ""
     MessagePack.unpack([0xd9, 0x01].pack('C*') + 'a').should == "a"