diff options
| -rw-r--r-- | lib/arel/nodes/node.rb | 7 | ||||
| -rw-r--r-- | lib/arel/visitors.rb | 1 | ||||
| -rw-r--r-- | lib/arel/visitors/depth_first.rb | 75 | ||||
| -rw-r--r-- | test/visitors/test_depth_first.rb | 76 |
4 files changed, 159 insertions, 0 deletions
diff --git a/lib/arel/nodes/node.rb b/lib/arel/nodes/node.rb index 841b954db9..90c63f0be9 100644 --- a/lib/arel/nodes/node.rb +++ b/lib/arel/nodes/node.rb @@ -3,6 +3,8 @@ module Arel ### # Abstract base class for all AST nodes class Node + include Enumerable + ### # Factory method to create a Nodes::Not node that has the recipient of # the caller as a child. @@ -32,6 +34,11 @@ module Arel viz = Visitors.for engine viz.accept self end + + # Iterate through AST, nodes will be yielded depth-first + def each &block + Visitors::DepthFirst.new(block).accept self + end end end end diff --git a/lib/arel/visitors.rb b/lib/arel/visitors.rb index 5d4c6084b2..2b0a06d534 100644 --- a/lib/arel/visitors.rb +++ b/lib/arel/visitors.rb @@ -1,4 +1,5 @@ require 'arel/visitors/visitor' +require 'arel/visitors/depth_first' require 'arel/visitors/to_sql' require 'arel/visitors/sqlite' require 'arel/visitors/postgresql' diff --git a/lib/arel/visitors/depth_first.rb b/lib/arel/visitors/depth_first.rb new file mode 100644 index 0000000000..19796b6e72 --- /dev/null +++ b/lib/arel/visitors/depth_first.rb @@ -0,0 +1,75 @@ +module Arel + module Visitors + class DepthFirst < Arel::Visitors::Visitor + def initialize block = nil + @block = block || Proc.new + end + + private + + def binary o + visit o.left + visit o.right + @block.call o + end + alias :visit_Arel_Nodes_And :binary + alias :visit_Arel_Nodes_Assignment :binary + alias :visit_Arel_Nodes_Between :binary + alias :visit_Arel_Nodes_DoesNotMatch :binary + alias :visit_Arel_Nodes_Equality :binary + alias :visit_Arel_Nodes_GreaterThan :binary + alias :visit_Arel_Nodes_GreaterThanOrEqual :binary + alias :visit_Arel_Nodes_In :binary + alias :visit_Arel_Nodes_LessThan :binary + alias :visit_Arel_Nodes_LessThanOrEqual :binary + alias :visit_Arel_Nodes_Matches :binary + alias :visit_Arel_Nodes_NotEqual :binary + alias :visit_Arel_Nodes_NotIn :binary + alias :visit_Arel_Nodes_Or :binary + + def visit_Arel_Attribute o + visit o.relation + visit o.name + @block.call o + end + alias :visit_Arel_Attributes_Integer :visit_Arel_Attribute + alias :visit_Arel_Attributes_Float :visit_Arel_Attribute + alias :visit_Arel_Attributes_String :visit_Arel_Attribute + alias :visit_Arel_Attributes_Time :visit_Arel_Attribute + alias :visit_Arel_Attributes_Boolean :visit_Arel_Attribute + alias :visit_Arel_Attributes_Attribute :visit_Arel_Attribute + + def visit_Arel_Table o + visit o.name + @block.call o + end + + def terminal o + @block.call o + end + alias :visit_Arel_Nodes_SqlLiteral :terminal + alias :visit_Arel_SqlLiteral :terminal + alias :visit_BigDecimal :terminal + alias :visit_Date :terminal + alias :visit_DateTime :terminal + alias :visit_FalseClass :terminal + alias :visit_Fixnum :terminal + alias :visit_Float :terminal + alias :visit_NilClass :terminal + alias :visit_String :terminal + alias :visit_Symbol :terminal + alias :visit_Time :terminal + alias :visit_TrueClass :terminal + + def visit_Array o + o.each { |i| visit i } + @block.call o + end + + def visit_Hash o + o.each { |k,v| visit(k); visit(v) } + @block.call o + end + end + end +end diff --git a/test/visitors/test_depth_first.rb b/test/visitors/test_depth_first.rb new file mode 100644 index 0000000000..eff60576fc --- /dev/null +++ b/test/visitors/test_depth_first.rb @@ -0,0 +1,76 @@ +require 'helper' + +module Arel + module Visitors + class TestDepthFirst < MiniTest::Unit::TestCase + Collector = Struct.new(:calls) do + def call object + calls << object + end + end + + def setup + @collector = Collector.new [] + @visitor = Visitors::DepthFirst.new @collector + end + + [ + Arel::Nodes::And, + Arel::Nodes::Assignment, + Arel::Nodes::Between, + Arel::Nodes::DoesNotMatch, + Arel::Nodes::Equality, + Arel::Nodes::GreaterThan, + Arel::Nodes::GreaterThanOrEqual, + Arel::Nodes::In, + Arel::Nodes::LessThan, + Arel::Nodes::LessThanOrEqual, + Arel::Nodes::Matches, + Arel::Nodes::NotEqual, + Arel::Nodes::NotIn, + Arel::Nodes::Or, + ].each do |klass| + define_method("test_#{klass.name.gsub('::', '_')}") do + binary = klass.new(:a, :b) + @visitor.accept binary + assert_equal [:a, :b, binary], @collector.calls + end + end + + [ + Arel::Attributes::Integer, + Arel::Attributes::Float, + Arel::Attributes::String, + Arel::Attributes::Time, + Arel::Attributes::Boolean, + Arel::Attributes::Attribute + ].each do |klass| + define_method("test_#{klass.name.gsub('::', '_')}") do + binary = klass.new(:a, :b) + @visitor.accept binary + assert_equal [:a, :b, binary], @collector.calls + end + end + + def test_table + relation = Arel::Table.new(:users) + @visitor.accept relation + assert_equal ['users', relation], @collector.calls + end + + def test_array + node = Nodes::Or.new(:a, :b) + list = [node] + @visitor.accept list + assert_equal [:a, :b, node, list], @collector.calls + end + + def test_hash + node = Nodes::Or.new(:a, :b) + hash = { node => node } + @visitor.accept hash + assert_equal [:a, :b, node, :a, :b, node, hash], @collector.calls + end + end + end +end |