def _test_opt(self, source, constants, expected_source=None, expected_new_bindings=None, print_source=False):
""" Test that with given constants, optimized_ast transforms
source to expected_source.
It :expected_new_bindings: is given, we check that they
are among new bindings returned by optimizer.
"""
if print_source:
print ast_to_string(ast.parse(shift_source(source)))
if expected_source is None:
expected_source = source
ast_tree = ast.parse(shift_source(source))
new_ast, bindings = optimized_ast(ast_tree, constants)
self.assertASTEqual(new_ast, ast.parse(shift_source(expected_source)))
if expected_new_bindings:
for k in expected_new_bindings:
if k not in bindings:
print "bindings:", bindings
self.assertEqual(bindings[k], expected_new_bindings[k])
def generic_visit(self, node):
''' Completly substite parent class "generic_visit", in order to
be able to insert some code at the line before current expression
(e.g. when inlining functions).
Also do some logging.
'''
prefix = '--' * self._depth
logger.debug('%s visit:\n%s', prefix, ast_to_string(node))
self._depth += 1
# copy-paste from ast.py, added self._current_block handling
block_fields = ['body', 'orelse'] # TODO more?
for field, old_value in ast.iter_fields(node):
old_value = getattr(node, field, None)
if isinstance(old_value, list):
new_values = []
if field in block_fields:
parent_block = self._current_block
self._current_block = new_values
try:
for value in old_value:
if isinstance(value, ast.AST):
value = self.visit(value)
if value is None:
continue
elif not isinstance(value, ast.AST):
new_values.extend(value)
continue
new_values.append(value)
old_value[:] = new_values
finally: # restore self._current_block
if field in block_fields:
self._current_block = parent_block
elif isinstance(old_value, ast.AST):
new_node = self.visit(old_value)
if new_node is None:
delattr(node, field)
else:
setattr(node, field, new_node)
# end of copy-paste
self._depth -= 1
logger.debug('%s result:\n%s', prefix, ast_to_string(node))
return node
