def test_mutiple_returns(self): source = ''' def f(x, y, z='foo'): if x: b = y + list(x) return b else: return z ''' ast_tree = ast.parse(shift_source(source)) expected_source = ''' def f(__ast_pe_var_4, __ast_pe_var_5, __ast_pe_var_6='foo'): if __ast_pe_var_4: __ast_pe_var_7 = __ast_pe_var_5 + list(__ast_pe_var_4) __ast_pe_var_8 = __ast_pe_var_7 break else: __ast_pe_var_8 = __ast_pe_var_6 break ''' inliner = Inliner(3, get_locals(ast_tree)) new_ast = inliner.visit(ast_tree) self.assertASTEqual(new_ast, ast.parse(shift_source(expected_source))) self.assertEqual(inliner.get_var_count(), 8) self.assertEqual(inliner.get_return_var(), '__ast_pe_var_8') self.assertEqual(inliner.get_bindings(), { 'x': '__ast_pe_var_4', 'y': '__ast_pe_var_5', 'z': '__ast_pe_var_6', 'b': '__ast_pe_var_7'})
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])