lib/treevisitor/tree_node.rb in gf-treevisitor-0.0.18 vs lib/treevisitor/tree_node.rb in gf-treevisitor-0.0.19

@@ -33,18 +33,38 @@
     def node(*args, &block)
       parent_node = @scope_stack.length > 0 ? @scope_stack[-1] : nil
       args << parent_node
       tree_node = @tree_node_class.new(*args)
       @scope_stack.push tree_node
-      class_eval(&block) if block 
+      if block
+        if block.arity == 0 || block.arity == -1
+          class_eval(&block)
+        elsif block.arity == 1
+          new_block = Proc.new{ block.call(tree_node) }
+          class_eval(&new_block)
+        else
+          raise "block take too much arguments #{block.arity}"
+        end
+      end
       @scope_stack.pop
     end
 
-    def leaf(*args)
+    def leaf(*args, &block)
       tree_node = @scope_stack[-1]
       args << tree_node
-      @leaf_node_class.new(*args)
+      leaf_node = @leaf_node_class.new(*args)
+      if block
+        if block.arity == 0 || block.arity == -1
+          class_eval(&block)
+        elsif block.arity == 1
+          new_block = Proc.new{ block.call(leaf_node) }
+          class_eval(&new_block)
+        else
+          raise "block take too much arguments #{block.arity}"
+        end
+      end
+      leaf_node
     end
   end
 
   attr_reader :leaves
   attr_reader :children
@@ -87,11 +107,11 @@
         
   def add_leaf( leaf )
     return if leaf.parent == self
     if not leaf.parent.nil?
       leaf.remove_from_parent
-    end  
+    end
     leaf.parent = self
     if @leaves.length > 0
       @leaves.last.next = leaf
       leaf.prev = @leaves.last
     else
@@ -106,11 +126,11 @@
     return if tree_node.parent == self
     if not tree_node.parent.nil?
       tree_node.remove_from_parent
     else
       tree_node.prefix_path = nil
-    end  
+    end
     tree_node.invalidate
     tree_node.parent = self
     if @children.length > 0
       @children.last.next = tree_node
       tree_node.prev = @children.last
@@ -155,11 +175,11 @@
     str = ""
 
     if root?
       str << to_s << "\n"
     else
-      str << prefix 
-      if self.next 
+      str << prefix
+      if self.next
         str << "|-- "
       else
         str << "\`-- "
       end
       str << to_s << "\n"