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')
