lib/better_html/ast/iterator.rb in better_html-1.0.1 vs lib/better_html/ast/iterator.rb in better_html-1.0.2

@@ -3,26 +3,26 @@
 
 module BetterHtml
   module AST
     class Iterator
       def initialize(types, &block)
-        @types = Array.wrap(types)
+        @types = types.nil? ? nil : Array.wrap(types)
         @block = block
       end
 
       def traverse(node)
         return unless node.is_a?(::AST::Node)
-        @block.call(node) if @types.include?(node.type)
+        @block.call(node) if @types.nil? || @types.include?(node.type)
         traverse_all(node)
       end
 
       def traverse_all(nodes)
         nodes.to_a.each do |node|
           traverse(node) if node.is_a?(::AST::Node)
         end
       end
 
-      def self.descendants(root_node, type, &block)
+      def self.descendants(root_node, type)
         Enumerator.new do |yielder|
           t = new(type) { |node| yielder << node }
           t.traverse(root_node)
         end
       end