def test_rename_kwarg(self): # expected expected = ''' def test3(x, y): data4 = paddle.tensor.argmax(input=x, axis=2, out=None, dtype=None, keepdims=False, name=None) ''' # captured input input_test = ''' def test3(x, y): data4 = paddle.fluid.layers.argmax(x=x, axis=2) ''' # get captured root = gast.parse(input_test) root_expected = gast.parse(expected) print("*******test_3_source_code*******") # print(inspect.getsource(test1)) root = transformer(root, modify_dict) print("*******test_rename_kwarg*******") print(astor.to_source(gast.gast_to_ast(root))) captured = astor.to_source(gast.gast_to_ast(root)) expected = astor.to_source(gast.gast_to_ast(root_expected)) # testting self.assertEqual(expected, captured)
def test_add_kwarg(self): # expected expected = ''' def test1(): data4 = paddle.tensor.zeros(x=x, shape=[3, 2], dtype='float32', out=None, device=None) ''' # captured input input_test = ''' def test1(): data4 = paddle.fluid.layers.zeros(x=x, shape=[3, 2], dtype='float32', force_cpu=True) ''' # get captured # root = gast.parse(inspect.getsource(input_test)) root = gast.parse(input_test) root_expected = gast.parse(expected) print("*******test_1_source_code*******") # print(inspect.getsource(test1)) root = transformer(root, modify_dict) print("*******test_add_kwarg*******") print(astor.to_source(gast.gast_to_ast(root))) captured = astor.to_source(gast.gast_to_ast(root)) expected = astor.to_source(gast.gast_to_ast(root_expected)) # testting self.assertEqual(expected, captured)
def test_delete_kwarg(self): # expected expected = ''' def test2(x, y): data4 = paddle.tensor.add(x, y, name=None, out=None, alpha=1) ''' # captured input input_test = ''' def test2(x, y): data4 = paddle.fluid.layers.elementwise_add(x, y, axis=-1, act=None, name=None) ''' # get captured root = gast.parse(input_test) root_expected = gast.parse(expected) print("*******test_2_source_code*******") # print(inspect.getsource(test1)) root = transformer(root, modify_dict) print("*******test_delete_kwarg*******") print(astor.to_source(gast.gast_to_ast(root))) captured = astor.to_source(gast.gast_to_ast(root)) expected = astor.to_source(gast.gast_to_ast(root_expected)) # testting self.assertEqual(expected, captured)
def test_replace_full_name_transformer(self, import_map=import_dict): # expected expected = ''' import paddle import paddle.fluid as fluid import numpy as np from paddle.fluid.dygraph.nn import Conv2D, Pool2D, Linear class LeNet(fluid.dygraph.Layer): def __init__(self, name_scope, num_classes=1): super(LeNet, self).__init__(name_scope) self.conv1 = paddle.fluid.dygraph.nn.Conv2D(num_channels=1, num_filters=6, filter_size=5, act='sigmoid') self.pool1 = paddle.fluid.dygraph.nn.Pool2D(pool_size=2, pool_stride=2, pool_type='max') self.conv2 = paddle.fluid.dygraph.nn.Conv2D(num_channels=6, num_filters=16, filter_size=5, act='sigmoid') self.pool2 = paddle.fluid.dygraph.nn.Pool2D(pool_size=2, pool_stride=2, pool_type='max') self.conv3 = paddle.fluid.dygraph.nn.Conv2D(num_channels=16, num_filters=120, filter_size=4, act='sigmoid') self.fc1 = paddle.fluid.dygraph.nn.Linear(input_dim=120, output_dim=64, act='sigmoid') self.fc2 = paddle.fluid.dygraph.nn.Linear(input_dim=64, output_dim=num_classes) def forward(self, x): x = self.conv1(x) x = self.pool1(x) x = self.conv2(x) x = self.pool2(x) x = self.conv3(x) x = paddle.fluid.layers.reshape(x, [x.shape[0], -1]) x = self.fc1(x) x = self.fc2(x) return x ''' # captured input input_test = gast.parse(inspect.getsource(dygraph_lenet)) print("*******original code*******") original_code = astor.to_source(gast.gast_to_ast(input_test)) original_code_formated = FormatCode(original_code) print(original_code_formated[0]) # get captured root = gast.parse(input_test) root_expected = gast.parse(expected) print("*******start_transfer*******") root = transfomer_replace_test(root, import_map) print("*******code_for_captured_formated*******") # print(astor.to_source(gast.gast_to_ast(root))) captured = astor.to_source(gast.gast_to_ast(root)) captured_formated = FormatCode(captured) print(captured_formated[0]) expected = astor.to_source(gast.gast_to_ast(root_expected)) expected_formated = FormatCode(expected) print("*******code_for_expected_formated*******") print(expected_formated[0]) # testting self.assertEqual(expected_formated, captured_formated)
def visit_Assert(self, node): if not self.static_analysis_visitor.is_tensor_node(node.test): return node cast_node = gast.Call( func=gast.parse("fluid.layers.cast").body[0].value, args=[node.test, gast.Constant(value="bool", kind=None)], keywords=[]) assert_node = gast.Call( func=gast.parse("fluid.layers.Assert").body[0].value, args=[cast_node], keywords=[]) return gast.Expr(value=assert_node)
def assertAstMatches(self, actual_node, expected_node_src, expr=True): if expr: # Ensure multi-line expressions parse. expected_node = gast.parse('({})'.format(expected_node_src)).body[0] expected_node = expected_node.value else: expected_node = gast.parse(expected_node_src).body[0] msg = 'AST did not match expected:\n{}\nActual:\n{}'.format( pretty_printer.fmt(expected_node), pretty_printer.fmt(actual_node)) self.assertTrue(ast_util.matches(actual_node, expected_node), msg)
def test_fusion_defDes(self): # expected expected = ''' def test1(x): data1 = paddle.nn.data(x=x) data2 = paddle.tensor.reshape(x=data1, shape=[1, 2, 3]) data4 = paddle.tensor.zeros(shape=[3, 2], dtype='float32', out=None, device=None) return abc def test2(x): data1 = paddle.tensor.tanh(input=x, out=None) data2 = paddle.tensor.zeros(shape=[3, 2], dtype='float32', out=None, device=None) def test3(x): data1 = paddle.nn.Sigmoid(x=x) data2 = paddle.tensor.sum(x, axis=0) def test4(x): data1 = paddle.tensor.max(input, dim=None, keep_dim=False, out=None, name=None) data2 = paddle.tensor.stack(x, dim=0, out=None) ''' # captured input input_test = ''' def test1(x): data1 = paddle.fluid.layers.data(x=x) data2 = paddle.fluid.layers.reshape(x=data1, shape=[1, 2, 3]) data3 = paddle.fluid.layers.zeros(shape=[3, 2], dtype='float32') data4 = paddle.fluid.layers.zeros(shape=[3, 2], dtype='float32', force_cpu=True) return abc def test2(x): data1 = paddle.fluid.layers.tanh(x=x) data2 = paddle.fluid.layers.zeros(shape=[3, 2], dtype='float32') def test3(x): data1 = paddle.fluid.layers.sigmoid(x=x) data2 = paddle.fluid.layers.reduce_sum(x, dim=0) def test4(x): data1 = paddle.fluid.layers.reduce_max(x=x) data2 = fluid.layers.stack(x, dim=0) ''' # get captured root = gast.parse(input_test) root_expected = gast.parse(expected) print("*******test_4_source_code*******") # print(inspect.getsource(test1)) root = transformer(root, modify_dict) print("*******test_fusion_function_kwarg*******") print(astor.to_source(gast.gast_to_ast(root))) captured = astor.to_source(gast.gast_to_ast(root)) expected = astor.to_source(gast.gast_to_ast(root_expected)) # testting self.assertEqual(expected, captured)
def parse_str(src): """Returns the AST of given piece of code.""" # TODO(mdan): This should exclude the module things are autowrapped in. if six.PY2 and re.search('\\Wprint\\s*\\(', src): # This special treatment is required because gast.parse is not aware of # whether print_function was present in the original context. src = 'from __future__ import print_function\n' + src parsed_module = gast.parse(src) parsed_module.body = parsed_module.body[1:] else: parsed_module = gast.parse(src) return parsed_module
def test_single_model_file(self): parser = argparse.ArgumentParser("Paddle API upgrade") parser.add_argument("--modify_dict", type=str, default="../dict/modify.dict") parser.add_argument("--input", type=str, default='class_cnn.py') parser.add_argument("--output", type=str, default='class_cnn_update.py') args = parser.parse_args() # read captured with open('class_cnn.py', 'r') as fp: data_captured = fp.read() with open('class_cnn_expected.py', 'r') as fp: data_expected = fp.read() # expected expected = data_expected # captured input input_test = data_captured # # get captured root = gast.parse(input_test) root_expected = gast.parse(expected) print("*******test_5_source_code*******") root = transformer(root, modify_dict) print("*******test_model_single_file*******") print(astor.to_source(gast.gast_to_ast(root))) captured = astor.to_source(gast.gast_to_ast(root)) captured_formated = FormatCode(captured) with open('class_cnn_update.py', 'w') as fp: fp.write(captured_formated[0]) expected = astor.to_source(gast.gast_to_ast(root_expected)) data_expected_formated = FormatCode(captured) # read captured with open('class_cnn_update.py', 'r') as fp: data_captured = fp.read() captured = astor.to_source(gast.gast_to_ast(root)) data_captured_formated = FormatCode(captured) # testting self.assertEqual(data_expected_formated, data_captured_formated)
def parse_str(src): """Returns the AST of given piece of code.""" # TODO(mdan): This should exclude the module things are autowrapped in. if six.PY2 and '.print(' in src: # This special treatment is required because gast.parse is not aware of # whether print_function was present in the original context. src = 'from __future__ import print_function\n' + src parsed_module = gast.parse(src) parsed_module.body = parsed_module.body[1:] else: parsed_module = gast.parse(src) return parsed_module
def assertAstMatches(self, actual_node, expected_node_src): expected_node = gast.parse(expected_node_src).body[0] msg = 'AST did not match expected:\n{}\nActual:\n{}'.format( pretty_printer.fmt(expected_node), pretty_printer.fmt(actual_node)) self.assertTrue(ast_util.matches(actual_node, expected_node), msg)
def pyfunc_to_ast(pyfunc): """ Transform py func to AST node """ src_code = textwrap.dedent(inspect.getsource(pyfunc)) ast_root = gast.parse(src_code) return ast_root
def visit_Attribute(self, node): self.generic_visit(node) attr_full_name = get_attribute_full_name(node) if attr_full_name == self.old_name: new_name_node = gast.parse(self.new_name).body[0].value return new_name_node return node
def test_paddle_api_with_andOr(self): code_or = """ def foo(x): if 2 > 1 and fluid.layers.shape(x)[0] > 16 or x is not None : x = x + 1 return x """ code_and = """ def foo(x): if 2 > 1 and fluid.layers.shape(x)[0] > 16 and x is not None : x = x + 1 return x """ for code in [code_or, code_and]: code = textwrap.dedent(code) node = gast.parse(code) static_analysis_visitor = StaticAnalysisVisitor(node) test_node = node.body[0].body[0].test if_visitor = IfConditionVisitor(test_node, static_analysis_visitor) self.assertTrue(if_visitor.is_control_flow()) new_node, assign_nodes = if_visitor.transform() # Transformation result: # bool_tensor_0 = fluid.layers.fill_constant(shape=[1], dtype='bool', value=bool(2 > 1)) # bool_tensor_1 = fluid.layers.fill_constant(shape=[1], dtype='bool', value=bool(x is not None)) # logic_and_0 = fluid.layers.logical_and(x=bool_tensor_0, y=fluid.layers.shape(x)[0] > 16) # logic_and_1 = fluid.layers.logical_and(x=logic_and_0, y=bool_tensor_1) for code_and # logic_or_0= fluid.layers.logical_or(x=logic_and_0, y=bool_tensor_1) for code_and self.assertTrue(isinstance(new_node, gast.Name)) self.assertTrue(len(assign_nodes) == 4)
def test_ast_graph_nodes(self): """Check node IDs, node types, and forward mapping.""" root = gast.parse( textwrap.dedent("""\ pass def foo(n): if n <= 1: return 1 """)) graph, forward_map = py_ast_graphs.py_ast_to_graph(root) # pytype: disable=attribute-error self.assertIn("root__Module", graph) self.assertEqual(graph["root__Module"].node_type, "Module") self.assertEqual(forward_map[id(root)], "root__Module") self.assertIn("root_body_1__Module_body-seq-helper", graph) self.assertEqual( graph["root_body_1__Module_body-seq-helper"].node_type, "Module_body-seq-helper") self.assertIn("root_body_1_item_body_0_item__If", graph) self.assertEqual(graph["root_body_1_item_body_0_item__If"].node_type, "If") self.assertEqual(forward_map[id(root.body[1].body[0])], "root_body_1_item_body_0_item__If") self.assertIn("root_body_1_item_body_0_item_test_left__Name", graph) self.assertEqual( graph["root_body_1_item_body_0_item_test_left__Name"].node_type, "Name") self.assertEqual(forward_map[id(root.body[1].body[0].test.left)], "root_body_1_item_body_0_item_test_left__Name")
def __init__(self, *par_values): # Check Jupyter notebook input cache self.src = self.recover_source() par_names = self.__get_caller_params() parsed = ast.parse(self.src) capture = PCapture(parsed, par_names) self.graph = capture.graph
def test_ast_graph_optional_field_edges(self): """Test that edges for optional fields are correct.""" root = gast.parse("return 1\nreturn") graph, _ = py_ast_graphs.py_ast_to_graph(root) self.assertEqual( graph["root_body_0_item__Return"].out_edges["value_out"], [ graph_types.InputTaggedNode( node_id=graph_types.NodeId( "root_body_0_item_value__Constant"), in_edge=graph_types.InEdgeType("parent_in")) ]) self.assertEqual( graph["root_body_0_item_value__Constant"].out_edges["parent_out"], [ graph_types.InputTaggedNode( node_id=graph_types.NodeId("root_body_0_item__Return"), in_edge=graph_types.InEdgeType("value_in")) ]) self.assertEqual( graph["root_body_1_item__Return"].out_edges["value_out"], [ graph_types.InputTaggedNode( node_id=graph_types.NodeId("root_body_1_item__Return"), in_edge=graph_types.InEdgeType("value_missing")) ])
def test_ast_graph_conforms_to_schema(self): # Some example code using a few different syntactic constructs, to cover # a large set of nodes in the schema root = gast.parse( textwrap.dedent("""\ def foo(n): if n <= 1: return 1 else: return foo(n-1) + foo(n-2) def bar(m, n) -> int: x = n for i in range(m): if False: continue x = x + i while True: break return x x0 = 1 + 2 - 3 * 4 / 5 x1 = (1 == 2) and (3 < 4) and (5 > 6) x2 = (7 <= 8) and (9 >= 10) or (11 != 12) x2 = bar(13, 14 + 15) """)) graph, _ = py_ast_graphs.py_ast_to_graph(root) # Graph should match the schema schema_util.assert_conforms_to_schema(graph, py_ast_graphs.SCHEMA)
def _replace_after_node_to_if_in_stmt_list( self, stmt_list, node, return_name, parent_node_of_return): i = index_in_list(stmt_list, node) if i < 0 or i >= len(stmt_list): return False if i == len(stmt_list) - 1: # No need to add, we consider this as added successfully return True if_stmt = gast.If(test=gast.UnaryOp( op=gast.Not(), operand=gast.Name( id=return_name, ctx=gast.Store(), annotation=None, type_comment=None)), body=stmt_list[i + 1:], orelse=[]) stmt_list[i + 1:] = [if_stmt] # Here assume that the parent node of return is gast.If if isinstance(parent_node_of_return, gast.If): # Prepend control flow boolean nodes such as '__return@1 = False' node_str = "{} = paddle.jit.dy2static.create_bool_as_type({}, False)".format( return_name, ast_to_source_code(parent_node_of_return.test).strip()) assign_false_node = gast.parse(node_str).body[0] stmt_list[i:i] = [assign_false_node] return True
def create_convert_shape_node(var_shape_node, slice_node=None, in_control_flow=False): assert isinstance(var_shape_node, (gast.Attribute, gast.Subscript)) if isinstance(var_shape_node, gast.Attribute): args = [ast_to_source_code(var_shape_node.value).strip()] # (1) A slice can be a simple number such as 1, -2, i.e. gast.Index or gast.Constant # (2) A slice can also be represented by bounds such as 2:-1, i.e. not gast.Index or gast.Constant # In (1) case, we pass the number as 'idx' argument in convert_var_shape # In (2) case, we have to make it like `convert_var_shape(x)[slice]` if slice_node is not None and slice_is_num(slice_node): args.append(ast_to_source_code(slice_node.slice).strip()) convert_var_shape_func = "paddle.jit.dy2static.convert_var_shape({}, in_control_flow={})".format( ",".join(args), in_control_flow) api_shape_node = gast.parse(convert_var_shape_func).body[0].value if slice_node is not None and not slice_is_num(slice_node): return gast.Subscript(value=api_shape_node, slice=slice_node.slice, ctx=gast.Load()) return api_shape_node if isinstance(var_shape_node, gast.Subscript): result_node = copy.deepcopy(var_shape_node) result_node = create_convert_shape_node(result_node.value, result_node, in_control_flow) return result_node
def test_nested_loop_vars(self): func = self.nested_for_loop_func test_func = inspect.getsource(func) gast_root = gast.parse(test_func) name_visitor = NameVisitor(gast_root) self.loop_var_names = [ set(["j", "two"]), set(["i", "three", "b"]), set(["i", "j"]), ] self.create_var_names = [set(), set(["b"]), set()] i = 0 for node in gast.walk(gast_root): if isinstance(node, (gast.While, gast.For)): loop_var_names, create_var_names = name_visitor.get_loop_var_names( node) self.assertEqual( loop_var_names, self.loop_var_names[i], msg="loop_var_names : {}, \nexpected loop_var_names : {}". format(loop_var_names, self.loop_var_names[i])) self.assertEqual( create_var_names, self.create_var_names[i], msg= "i = {}\ncreate_var_names : {}, \nexpected create_var_names : {}" .format(i, create_var_names, self.create_var_names[i])) i += 1
def testCompile(self): srcs = glob.glob(os.path.join(gast.__path__[0], '*.py')) for src_py in srcs: with open(src_py) as f: content = f.read() gnode = gast.parse(content) compile(gast.gast_to_ast(gnode), src_py, 'exec')
def _replace_pop(self, node): """ Replace a pop statement for a list or dict. For example: list_a = [0,1,2,3,4] x = list_a.pop() # --> convert_pop(list_a) y = list_a.pop(1) # --> convert_pop(list_a, 1) dict_a = {"red":0, "blue":1, "yellow":2} m = dict_a.pop("red") # --> convert_pop(dict_a, "red") n = dict_a.pop("black", 3) # --> convert_pop(dict_a, "black", 3) """ assert isinstance(node, gast.Call) assert isinstance(node.func, gast.Attribute) target_node = node.func.value target_str = ast_to_source_code(target_node).strip() args_str = [ast_to_source_code(arg).strip() for arg in node.args] # NOTE(liym27): # 1. pop stmt for a list if len(args_str) == 0 # 2. pop stmt for a list or dict if len(args_str) == 1 # 3. pop stmt for a dict if len(args_str) == 2 if len(args_str) <= 2: new_pop_str = "paddle.jit.dy2static.convert_pop({}, {})"\ .format(target_str, ",".join(args_str)) new_pop_node = gast.parse(new_pop_str).body[0].value return new_pop_node else: return node
def test_get_name_ids(self): source = textwrap.dedent(self.source) root = gast.parse(source) all_name_ids = get_name_ids([root]) self.assertDictEqual( self.transfer_dict(self.all_name_ids), self.transfer_dict(all_name_ids))
def _create_bool_op_node(self, nodes, api_type): ''' NOTE(liym27): The arguments of function convert_logical_XX should be callable so that they can be run according to the actual order. In `convert_logical_and(lambda:x>1, lambda:y<1)`, `lambda:y<1` must be run after `lambda:x>1`, If `x>1` is False, `y<1` should NOT be run. ''' assert len( nodes ) > 1, "The length of BoolOp should be at least 2, but received {}.".format( len(nodes)) if len(nodes) > 2: # Creates logic_and/logic_or node recursively. pre_logic_node = self._create_bool_op_node(nodes[:2], api_type) if len(nodes[2:]) == 1: post_logic_node = nodes[2] else: post_logic_node = self._create_bool_op_node(nodes[2:], api_type) nodes = [pre_logic_node] + [post_logic_node] args = [ast_to_source_code(child) for child in nodes] new_node_str = "paddle.jit.dy2static.convert_logical_{}(lambda:{}, lambda:{})".format( api_type, args[0], args[1]) # NOTE: gast.parse return Module(body=[expr(...)]) new_node = gast.parse(new_node_str).body[0].value return new_node
def test_Bytes(self): code = 'b"0012"' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec') norm = ("Module(body=[Expr(value=Constant(value=b'0012', " "kind=None))], type_ignores=[])") self.assertEqual(gast.dump(tree), norm)
def create_while_node(condition_name, body_name, loop_var_names): while_args = [] while_args.append( gast.Name(id=condition_name, ctx=gast.Param(), annotation=None, type_comment=None)) while_args.append( gast.Name(id=body_name, ctx=gast.Param(), annotation=None, type_comment=None)) assign_targets = [ gast.Name(id=var_name, ctx=gast.Param(), annotation=None, type_comment=None) for var_name in loop_var_names ] while_args.append(gast.List(elts=assign_targets, ctx=gast.Param())) while_func_id = gast.parse('fluid.layers.while_loop').body[0].value while_node = gast.Call(func=while_func_id, args=while_args, keywords=[]) assign_node = gast.Assign( targets=[gast.Tuple(elts=assign_targets, ctx=gast.Store())], value=while_node) return assign_node
def test_TryFinally(self): code = 'try:pass\nfinally:pass' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec') norm = ("Module(body=[Try(body=[Pass()], handlers=[], orelse=[], " "finalbody=[Pass()])], type_ignores=[])") self.assertEqual(gast.dump(tree), norm)
def test_compute_same_identifier_edges(self): list_node = gast.parse("[x, x, x, y, y]").body[0].value ast_to_node_id = { id(list_node.elts[0]): "x0", id(list_node.elts[1]): "x1", id(list_node.elts[2]): "x2", id(list_node.elts[3]): "y0", id(list_node.elts[4]): "y1", } same_identifier_edges = graph_edge_util.compute_same_identifier_edges( list_node, ast_to_node_id) self.assertCountEqual(same_identifier_edges, [ ("x0", "x0", "EXTRA_SAME_IDENTIFIER"), ("x0", "x1", "EXTRA_SAME_IDENTIFIER"), ("x0", "x2", "EXTRA_SAME_IDENTIFIER"), ("x1", "x0", "EXTRA_SAME_IDENTIFIER"), ("x1", "x1", "EXTRA_SAME_IDENTIFIER"), ("x1", "x2", "EXTRA_SAME_IDENTIFIER"), ("x2", "x0", "EXTRA_SAME_IDENTIFIER"), ("x2", "x1", "EXTRA_SAME_IDENTIFIER"), ("x2", "x2", "EXTRA_SAME_IDENTIFIER"), ("y0", "y0", "EXTRA_SAME_IDENTIFIER"), ("y0", "y1", "EXTRA_SAME_IDENTIFIER"), ("y1", "y0", "EXTRA_SAME_IDENTIFIER"), ("y1", "y1", "EXTRA_SAME_IDENTIFIER"), ])
def test_compute_jumps_out_edges(self): tree = gast.parse( textwrap.dedent("""\ def foo(): # tree.body[0] return # tree.body[0].body[0] while True: # tree.body[0].body[1] break # tree.body[0].body[1].body[0] continue # tree.body[0].body[1].body[1] return # tree.body[0].body[1].body[2] while True: # tree.body[0].body[1].body[3] break # tree.body[0].body[1].body[3].body[0] return 4 # tree.body[0].body[1].body[3].body[1] """)) expected_type = graph_edge_util.JUMPS_OUT_OF_EDGE_TYPE expected_targets = [ (tree.body[0].body[0], tree.body[0], expected_type), (tree.body[0].body[1].body[0], tree.body[0].body[1], expected_type), (tree.body[0].body[1].body[1], tree.body[0].body[1], expected_type), (tree.body[0].body[1].body[2], tree.body[0], expected_type), (tree.body[0].body[1].body[3].body[0], tree.body[0].body[1].body[3], expected_type), (tree.body[0].body[1].body[3].body[1], tree.body[0], expected_type), (tree.body[0].body[1].body[3].body[1].value, tree.body[0], expected_type), ] # For this test, we pretend that the AST nodes are the node ids. targets = graph_edge_util.compute_jumps_out_edges( tree, {id(x): x for x in gast.walk(tree)}) self.assertCountEqual(targets, expected_targets)
def make_ast( target_ast_node_count, rng = None, distribution = ( python_numbers_control_flow.DATAFLOW_FNS_DISTRIBUTION) ): """Generates an AST for this task. Args: target_ast_node_count: How many nodes to put in the AST. rng: Random state to use. distribution: Sampling distribution to use when building the AST. May also be a callable that produces a distribution given a random state. Returns: AST of a generated program. """ def root_build(body): """Given a list of statements, puts them into a function in a module.""" return gast.Module( body=[ gast.FunctionDef( name="random_function", args=_make_arguments( python_numbers_control_flow.make_name("a"), python_numbers_control_flow.make_name("b")), body=body, decorator_list=[], returns=None, type_comment=None) ], type_ignores=[]) root_template = python_numbers_control_flow.ASTWithHoles( cost=5, holes=[ top_down_refinement.Hole( python_numbers_control_flow.ASTHoleType.STMTS_NONEMPTY, python_numbers_control_flow.ASTHoleMetadata(("a", "b"), True, False, 0)) ], build=root_build) if rng is None: rng = np.random.RandomState() if callable(distribution): distribution = distribution(rng) tree = top_down_refinement.top_down_construct( root_object=root_template, target_cost=target_ast_node_count, refinement_distribution=distribution, rng=rng) # Re-parse the tree so that it is valid. This is required for program graph # analysis to work. return gast.parse(astunparse.unparse(gast.gast_to_ast(tree)))
def test_walk(self): code = 'x + 1' tree = gast.parse(code, mode='eval') dump = gast.dump(tree) norm = ("Expression(body=BinOp(left=Name(id='x', ctx=Load(), " "annotation=None), op=Add(), right=Num(n=1)))") self.assertEqual(dump, norm) self.assertEqual(len(list(gast.walk(tree))), 6)
def test_dump(self): code = 'lambda x: x' tree = gast.parse(code, mode='eval') dump = gast.dump(tree) norm = ("Expression(body=Lambda(args=arguments(args=[Name(id='x', " "ctx=Param(), annotation=None)], vararg=None, kwonlyargs=[], " "kw_defaults=[], kwarg=None, defaults=[]), body=Name(id='x', " "ctx=Load(), annotation=None)))") self.assertEqual(dump, norm)
def test_Raise(self): codes = ('raise Exception', 'raise "Exception"', 'raise Exception, "err"', 'raise Exception("err")', 'raise E, V, T',) for code in codes: tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def __init__(self, name, module=None): """Parameters are the name for the module (mandatory), and the ast.Module node (optional) in the case the current module is the main one. This differentiation is needed to avoid mangling function name for functions defined in the main module. """ self.is_main_module = True self.node = module if self.node is None: # Not main module, parse now the imported module self.is_main_module = False imported_module = importlib.import_module(name) self.node = ast.parse(inspect.getsource(imported_module)) assert isinstance(self.node, ast.Module) # Recursively add filename information to all nodes, for debug msg add_filename_field(self.node, name + ".py") # Mangle function imported, unless it is the main module self.to_be_mangled = not self.is_main_module self.name = name # Functions defined in this module and imported by another one. # This dict is used at the end of the process to gather functions to be # prepend at the beginning of the main pythran module self.exported_functions = dict() self.dependent_modules = dict() # Top-level function declared in this module self.functions = dict() # Functions imported as "from somemodule import func as func_alias" self.imported_functions = dict() # Regular module import. Keys are alias and values are module names self.imported_modules = dict() # Collect top-level functions and imports for decl in self.node.body: if isinstance(decl, ast.FunctionDef): # regular functions self.functions[decl.name] = decl elif isinstance(decl, ast.Import): # Module import for alias in decl.names: asname = alias.asname or alias.name self.imported_modules[asname] = alias.name elif isinstance(decl, ast.ImportFrom): # Function import module_name = decl.module for alias in decl.names: func_name = alias.name asname = alias.asname or func_name self.imported_functions[asname] = (module_name, func_name, None) elif isinstance(decl, ast.Assign): # FIXME : We ignore import of globals pass else: raise PythranSyntaxError('Unpythranizable module: %s' % name)
def parse(pm, code): # hacky way to turn OpenMP comments into strings code = re.sub(r'(\s*)#\s*(omp\s[^\n]+)', r'\1"\2"', code) # front end ir = ast.parse(code) # parse openmp directive pm.apply(GatherOMPData, ir) # extract docstrings _, docstrings = pm.apply(ExtractDocStrings, ir) # Handle user-defined import pm.apply(HandleImport, ir) # avoid conflicts with cxx keywords _, renamings = pm.apply(NormalizeIdentifiers, ir) check_syntax(ir) return ir, renamings, docstrings
def parse_str(src, preamble_len=0, single_node=True): """Returns the AST of given piece of code. Args: src: Text preamble_len: Int, indicates leading nodes in the parsed AST which should be dropped. single_node: Bool, whether `src` is assumed to be represented by exactly one AST node. Returns: ast.AST """ module_node = gast.parse(src) nodes = module_node.body if preamble_len: nodes = nodes[preamble_len:] if single_node: if len(nodes) != 1: raise ValueError('expected exactly one node node, found {}'.format(nodes)) return nodes[0] return nodes
def test_literal_eval_code(self): code = "[1, 3]" tree = ast.parse(code, mode='eval') gtree = gast.parse(code, mode='eval') self.assertEqual(ast.literal_eval(tree), ast.literal_eval(tree))
def test_parse(self): code = ''' def foo(x=1, *args, **kwargs): return x + y +len(args) + len(kwargs) ''' tree = gast.parse(code)
def test_get_docstring(self): code = 'def foo(): "foo"' tree = gast.parse(code) func = tree.body[0] docs = gast.get_docstring(func) self.assertEqual(docs, "foo")
def test_With(self): code = 'with open("any"): pass' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def parse_str(src): """Return the AST of given piece of code.""" return gast.parse(textwrap.dedent(src))
def test_TryFinally(self): code = 'try:pass\nfinally:pass' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def test_Call(self): code = 'foo(x, y=1, *args, **kwargs)' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def test_node_equality(self): node_a = gast.parse('y = x').body[0] node_b = gast.parse('y = x').body[0] self.assertNotEqual(node_a, node_b)
def parse_str(src): """Return the AST of given piece of code.""" return gast.parse(src)
def test_KeywordOnlyArgument(self): code = 'def foo(*, x=1): pass' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def test_TryExceptNamed(self): code = 'try:pass\nexcept e as f:pass\nelse:pass' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def parse_str(src): """Returns the AST of given piece of code.""" # TODO(mdan): This should exclude the module things are autowrapped in. return gast.parse(src)
def testParse(self): for src_py in self.srcs: with open(src_py) as f: content = f.read() gast.parse(content)
def test_unparse(self): for src_py in self.srcs: with open(src_py) as f: content = f.read() gnode = gast.parse(content) astunparse.unparse(gast.gast_to_ast(gnode))
def test_ArgAnnotation(self): code = 'def foo(x:int): pass' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def test_FormattedValue(self): code = 'e = 1; f"{e}"' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def test_JoinedStr(self): code = 'e = 1; f"e = {e}"' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec')
def raw_parse(code): # hacky way to turn OpenMP comments into strings code = re.sub(r'(\s*)#\s*(omp\s[^\n]+)', r'\1"\2"', code) return ast.parse(code)
def test_keyword_argument(self): code = 'def foo(**a): pass' tree = gast.parse(code) compile(gast.gast_to_ast(tree), '<test>', 'exec') gast.dump(tree, include_attributes=True)
def testCompile(self): for src_py in self.srcs: with open(src_py) as f: content = f.read() gnode = gast.parse(content) compile(gast.gast_to_ast(gnode), src_py, 'exec')