Allow comparison and serialization with sub-classes.

lib/tree.rb

@@ -291,7 +291,7 @@ def marshal_load(dumped_tree_array)
         content     = Marshal.load(node_hash[:content])
 
         if parent_name
-          nodes[name] = current_node = Tree::TreeNode.new(name, content)
+          nodes[name] = current_node = self.class.new(name, content)
           nodes[parent_name].add current_node
         else
           # This is the root node, hence initialize self.
@@ -952,7 +952,7 @@ def previous_sibling
     #                   this node is a 'predecessor'. Returns 'nil' if the other
     #                   object is not a 'Tree::TreeNode'.
     def <=>(other)
-      return nil if other == nil || other.class != Tree::TreeNode
+      return nil if other == nil || !other.is_a?(Tree::TreeNode)
       self.name <=> other.name
     end
 

spec/tree_spec.rb

@@ -11,6 +11,7 @@
 require 'spec_helper'
 
 describe Tree do
+  class SpecializedTreeNode < Tree::TreeNode; end
 
   shared_examples_for 'any detached node' do
     it 'should not equal "Object.new"' do
@@ -71,4 +72,48 @@
 
     it_behaves_like 'any detached node'
   end
+
+  describe 'comparison' do
+    let(:node1) { Tree::TreeNode.new('a') }
+    let(:node2) { SpecializedTreeNode.new('b') }
+
+    it 'allows comparison of specialized tree nodes' do
+      expect(node1 <=> node2).to be_eql(-1)
+    end
+
+    it 'does not allow comparison with nil' do
+      expect(node1 <=> nil).to be(nil)
+    end
+
+    it 'does not allow comparison with other objects' do
+      expect(node1 <=> 'c').to be(nil)
+    end
+  end
+
+  describe 'serialization' do
+    let(:node1) { SpecializedTreeNode.new('a') }
+    let(:serialized_node1) { Marshal.dump(node1) }
+    let(:node2) { Tree::TreeNode.new('b') }
+    let(:serialized_node2) { Marshal.dump(node2) }
+    let(:tree) do
+      SpecializedTreeNode.new('root').tap do |root|
+        root << node1
+        root << node2
+      end
+    end
+    let(:serialized_tree) { Marshal.dump(tree) }
+
+    it 'parses the serialized specialized tree node correctly (root)' do
+      expect(Marshal.load(serialized_tree)).to be_a(SpecializedTreeNode)
+    end
+
+    it 'parses the serialized specialized tree node correctly (child)' do
+      expect(Marshal.load(serialized_tree).children.first).to \
+        be_a(SpecializedTreeNode)
+    end
+
+    it 'parses the serialized tree node correctly' do
+      expect(Marshal.load(serialized_node2)).to be_a(Tree::TreeNode)
+    end
+  end
 end