def test_visit_power(self):
node = PowerNode(sum_node_for_num(20), sum_node_for_num(5))
traverse_pre(node, self.visitor)
self.assertEqual(self.visitor.counts, [
'power', 'sum', 'mult', 'power', 'func_', 'number_20', 'sum',
'mult', 'power', 'func_', 'number_5'
])
def test_visit_sum(self):
node = SumNode(mult_node_for_num(30),
SumSub(SignNode('+'), sum_node_for_num(50)))
traverse_pre(node, self.visitor)
self.assertEqual(self.visitor.counts, [
'sum', 'mult', 'power', 'func_', 'number_30', 'sum-sub', 'sign',
'sum', 'mult', 'power', 'func_', 'number_50'
])
def test_visit_mult(self):
node = MultNode(power_node_for_num(20))
traverse_pre(node, self.visitor)
self.assertEqual(self.visitor.counts,
['mult', 'power', 'func_', 'number_20'])
def test_visit_name(self):
node = NameNode('id')
traverse_pre(node, self.visitor)
self.assertEqual(self.visitor.counts, ['name'])
def test_visit_funcall(self):
node = FuncallNode(NameNode('sin'), sum_node_for_num(30))
traverse_pre(node, self.visitor)
self.assertEqual(
self.visitor.counts,
['func_sin', 'name', 'sum', 'mult', 'power', 'func_', 'number_30'])
def test_visit_number(self):
node = NumberNode(20)
traverse_pre(node, self.visitor)
self.assertEqual(self.visitor.counts, ['number_20'])
def test_skip_power(self):
visitor = SkipPowerNodes()
node = PowerNode(funcall(20), expr(10, [mult_sub('*', 3)]))
traverse_pre(node, visitor)
self.assertEqual(visitor.counts, ['power', 'func_', 'number_20'])
def test_stop_at_root(self):
visitor = HeadSequenceVisitor()
node = mult_node_for_num(30)
traverse_pre(node, visitor)
print(visitor.counts)
self.assertEqual(visitor.counts, ['mult'])
