THRIFT-697. Union support in Ruby

git-svn-id: https://svn.apache.org/repos/asf/incubator/thrift/trunk@910700 13f79535-47bb-0310-9956-ffa450edef68

6 files changed, 265 insertions, 107 deletions

diff --git a/lib/rb/lib/thrift.rb b/lib/rb/lib/thrift.rb
index 4d4e130ae..02d67b8ba 100644
--- a/lib/rb/lib/thrift.rb
+++ b/lib/rb/lib/thrift.rb
@@ -28,6 +28,8 @@ require 'thrift/types'
 require 'thrift/processor'
 require 'thrift/client'
 require 'thrift/struct'
+require 'thrift/union'
+require 'thrift/struct_union'
 
 # serializer
 require 'thrift/serializer/serializer'
diff --git a/lib/rb/lib/thrift/protocol/binary_protocol_accelerated.rb b/lib/rb/lib/thrift/protocol/binary_protocol_accelerated.rb
index eaf64f6be..70ea652c8 100644
--- a/lib/rb/lib/thrift/protocol/binary_protocol_accelerated.rb
+++ b/lib/rb/lib/thrift/protocol/binary_protocol_accelerated.rb
@@ -29,7 +29,11 @@ See MemoryBuffer and BufferedTransport for examples.
 module Thrift
   class BinaryProtocolAcceleratedFactory < BaseProtocolFactory
     def get_protocol(trans)
-      BinaryProtocolAccelerated.new(trans)
+      if (defined? BinaryProtocolAccelerated)
+        BinaryProtocolAccelerated.new(trans)
+      else
+        BinaryProtocol.new(trans)
+      end
     end
   end
 end
diff --git a/lib/rb/lib/thrift/struct.rb b/lib/rb/lib/thrift/struct.rb
index dfc8a2fec..9e52073b3 100644
--- a/lib/rb/lib/thrift/struct.rb
+++ b/lib/rb/lib/thrift/struct.rb
@@ -65,26 +65,7 @@ module Thrift
       end
       fields_with_default_values
     end
-
-    def name_to_id(name)
-      names_to_ids = self.class.instance_variable_get("@names_to_ids")
-      unless names_to_ids
-        names_to_ids = {}
-        struct_fields.each do |fid, field_def|
-          names_to_ids[field_def[:name]] = fid
-        end
-        self.class.instance_variable_set("@names_to_ids", names_to_ids)
-      end
-      names_to_ids[name]
-    end
-
-    def each_field
-      struct_fields.keys.sort.each do |fid|
-        data = struct_fields[fid]
-        yield fid, data
-      end
-    end
-
+    
     def inspect(skip_optional_nulls = true)
       fields = []
       each_field do |fid, field_info|
@@ -115,7 +96,8 @@ module Thrift
       each_field do |fid, field_info|
         name = field_info[:name]
         type = field_info[:type]
-        if (value = instance_variable_get("@#{name}"))
+        value = instance_variable_get("@#{name}")
+        unless value.nil?
           if is_container? type
             oprot.write_field_begin(name, type, fid)
             write_container(oprot, value, field_info)
@@ -210,89 +192,5 @@ module Thrift
         iprot.skip(ftype)
       end
     end
-
-    def read_field(iprot, field = {})
-      case field[:type]
-      when Types::STRUCT
-        value = field[:class].new
-        value.read(iprot)
-      when Types::MAP
-        key_type, val_type, size = iprot.read_map_begin
-        value = {}
-        size.times do
-          k = read_field(iprot, field_info(field[:key]))
-          v = read_field(iprot, field_info(field[:value]))
-          value[k] = v
-        end
-        iprot.read_map_end
-      when Types::LIST
-        e_type, size = iprot.read_list_begin
-        value = Array.new(size) do |n|
-          read_field(iprot, field_info(field[:element]))
-        end
-        iprot.read_list_end
-      when Types::SET
-        e_type, size = iprot.read_set_begin
-        value = Set.new
-        size.times do
-          element = read_field(iprot, field_info(field[:element]))
-          value << element
-        end
-        iprot.read_set_end
-      else
-        value = iprot.read_type(field[:type])
-      end
-      value
-    end
-
-    def write_data(oprot, value, field)
-      if is_container? field[:type]
-        write_container(oprot, value, field)
-      else
-        oprot.write_type(field[:type], value)
-      end
-    end
-
-    def write_container(oprot, value, field = {})
-      case field[:type]
-      when Types::MAP
-        oprot.write_map_begin(field[:key][:type], field[:value][:type], value.size)
-        value.each do |k, v|
-          write_data(oprot, k, field[:key])
-          write_data(oprot, v, field[:value])
-        end
-        oprot.write_map_end
-      when Types::LIST
-        oprot.write_list_begin(field[:element][:type], value.size)
-        value.each do |elem|
-          write_data(oprot, elem, field[:element])
-        end
-        oprot.write_list_end
-      when Types::SET
-        oprot.write_set_begin(field[:element][:type], value.size)
-        value.each do |v,| # the , is to preserve compatibility with the old Hash-style sets
-          write_data(oprot, v, field[:element])
-        end
-        oprot.write_set_end
-      else
-        raise "Not a container type: #{field[:type]}"
-      end
-    end
-
-    CONTAINER_TYPES = []
-    CONTAINER_TYPES[Types::LIST] = true
-    CONTAINER_TYPES[Types::MAP] = true
-    CONTAINER_TYPES[Types::SET] = true
-    def is_container?(type)
-      CONTAINER_TYPES[type]
-    end
-
-    def field_info(field)
-      { :type => field[:type],
-        :class => field[:class],
-        :key => field[:key],
-        :value => field[:value],
-        :element => field[:element] }
-    end
   end
 end
diff --git a/lib/rb/lib/thrift/struct_union.rb b/lib/rb/lib/thrift/struct_union.rb
new file mode 100644
index 000000000..9a5903f11
--- /dev/null
+++ b/lib/rb/lib/thrift/struct_union.rb
@@ -0,0 +1,126 @@
+# 
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+# 
+#   http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+require 'set'
+
+module Thrift
+  module Struct_Union
+    def name_to_id(name)
+      names_to_ids = self.class.instance_variable_get("@names_to_ids")
+      unless names_to_ids
+        names_to_ids = {}
+        struct_fields.each do |fid, field_def|
+          names_to_ids[field_def[:name]] = fid
+        end
+        self.class.instance_variable_set("@names_to_ids", names_to_ids)
+      end
+      names_to_ids[name]
+    end
+
+    def each_field
+      struct_fields.keys.sort.each do |fid|
+        data = struct_fields[fid]
+        yield fid, data
+      end
+    end
+
+    def read_field(iprot, field = {})
+      case field[:type]
+      when Types::STRUCT
+        value = field[:class].new
+        value.read(iprot)
+      when Types::MAP
+        key_type, val_type, size = iprot.read_map_begin
+        value = {}
+        size.times do
+          k = read_field(iprot, field_info(field[:key]))
+          v = read_field(iprot, field_info(field[:value]))
+          value[k] = v
+        end
+        iprot.read_map_end
+      when Types::LIST
+        e_type, size = iprot.read_list_begin
+        value = Array.new(size) do |n|
+          read_field(iprot, field_info(field[:element]))
+        end
+        iprot.read_list_end
+      when Types::SET
+        e_type, size = iprot.read_set_begin
+        value = Set.new
+        size.times do
+          element = read_field(iprot, field_info(field[:element]))
+          value << element
+        end
+        iprot.read_set_end
+      else
+        value = iprot.read_type(field[:type])
+      end
+      value
+    end
+
+    def write_data(oprot, value, field)
+      if is_container? field[:type]
+        write_container(oprot, value, field)
+      else
+        oprot.write_type(field[:type], value)
+      end
+    end
+
+    def write_container(oprot, value, field = {})
+      case field[:type]
+      when Types::MAP
+        oprot.write_map_begin(field[:key][:type], field[:value][:type], value.size)
+        value.each do |k, v|
+          write_data(oprot, k, field[:key])
+          write_data(oprot, v, field[:value])
+        end
+        oprot.write_map_end
+      when Types::LIST
+        oprot.write_list_begin(field[:element][:type], value.size)
+        value.each do |elem|
+          write_data(oprot, elem, field[:element])
+        end
+        oprot.write_list_end
+      when Types::SET
+        oprot.write_set_begin(field[:element][:type], value.size)
+        value.each do |v,| # the , is to preserve compatibility with the old Hash-style sets
+          write_data(oprot, v, field[:element])
+        end
+        oprot.write_set_end
+      else
+        raise "Not a container type: #{field[:type]}"
+      end
+    end
+
+    CONTAINER_TYPES = []
+    CONTAINER_TYPES[Types::LIST] = true
+    CONTAINER_TYPES[Types::MAP] = true
+    CONTAINER_TYPES[Types::SET] = true
+    def is_container?(type)
+      CONTAINER_TYPES[type]
+    end
+
+    def field_info(field)
+      { :type => field[:type],
+        :class => field[:class],
+        :key => field[:key],
+        :value => field[:value],
+        :element => field[:element] }
+    end
+  end
+end
\ No newline at end of file
diff --git a/lib/rb/lib/thrift/types.rb b/lib/rb/lib/thrift/types.rb
index 20e4ca2c1..cac52691a 100644
--- a/lib/rb/lib/thrift/types.rb
+++ b/lib/rb/lib/thrift/types.rb
@@ -57,7 +57,7 @@ module Thrift
               when Types::STRING
                 String
               when Types::STRUCT
-                Struct
+                [Struct, Union]
               when Types::MAP
                 Hash
               when Types::SET
diff --git a/lib/rb/lib/thrift/union.rb b/lib/rb/lib/thrift/union.rb
new file mode 100644
index 000000000..0b41ed495
--- /dev/null
+++ b/lib/rb/lib/thrift/union.rb
@@ -0,0 +1,128 @@
+# 
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+# 
+#   http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# 
+
+module Thrift
+  class Union
+    def initialize(name=nil, value=nil)
+      if name
+        if value.nil?
+          raise Exception, "Union #{self.class} cannot be instantiated with setfield and nil value!"
+        end
+
+        Thrift.check_type(value, struct_fields[name_to_id(name.to_s)], name) if Thrift.type_checking
+      elsif !value.nil?
+        raise Exception, "Value provided, but no name!"
+      end
+      @setfield = name
+      @value = value
+    end
+
+    def inspect
+      "<#{self.class} #{@setfield}: #{@value}>"
+    end
+
+    def read(iprot)
+      iprot.read_struct_begin
+      fname, ftype, fid = iprot.read_field_begin
+      handle_message(iprot, fid, ftype)
+      iprot.read_field_end
+
+      fname, ftype, fid = iprot.read_field_begin
+      raise "Too many fields for union" unless (ftype == Types::STOP) 
+
+      iprot.read_struct_end
+      validate
+    end
+
+    def write(oprot)
+      validate
+      oprot.write_struct_begin(self.class.name)
+
+      fid = self.name_to_id(@setfield.to_s)
+
+      field_info = struct_fields[fid]
+      type = field_info[:type]
+      if is_container? type
+        oprot.write_field_begin(@setfield, type, fid)
+        write_container(oprot, @value, field_info)
+        oprot.write_field_end
+      else
+        oprot.write_field(@setfield, type, fid, @value)
+      end
+
+      oprot.write_field_stop
+      oprot.write_struct_end
+    end
+
+    def ==(other)
+      other != nil && @setfield == other.get_set_field && @value == other.get_value
+    end
+
+    def eql?(other)
+      self.class == other.class && self == other
+    end
+
+    def hash
+      [self.class.name, @setfield, @value].hash
+    end
+
+    def self.field_accessor(klass, *fields)
+      fields.each do |field|
+        klass.send :define_method, "#{field}" do
+          if field == @setfield
+            @value
+          else 
+            raise RuntimeError, "#{field} is not union's set field."
+          end
+        end
+
+        klass.send :define_method, "#{field}=" do |value|
+          Thrift.check_type(value, klass::FIELDS.values.find {|f| f[:name].to_s == field.to_s }, field) if Thrift.type_checking
+          @setfield = field
+          @value = value
+        end
+      end
+    end
+
+    # get the symbol that indicates what the currently set field type is. 
+    def get_set_field
+      @setfield
+    end
+
+    # get the current value of this union, regardless of what the set field is.
+    # generally, you should only use this method when you don't know in advance
+    # what field to expect.
+    def get_value
+      @value
+    end
+
+    protected
+
+    def handle_message(iprot, fid, ftype)
+      field = struct_fields[fid]
+      if field and field[:type] == ftype
+        @value = read_field(iprot, field)
+        name = field[:name].to_sym
+        @setfield = name
+      else
+        iprot.skip(ftype)
+      end
+    end
+  end
+end
\ No newline at end of file