def finish_tree(self, tree: Node, filename: str) -> None:
# TODO: what about name clash between dotted-path imports and
# introspected locals?
imports_dict = get_import_lines(self.threadlocals.strategy_to_names)
insert_pos = _find_import_pos(tree)
def _sort_key(val):
left, right = val
left = left or ""
return (left, right)
sorted_tuples = sorted(
imports_dict.items(),
key=_sort_key,
reverse=True, # because we insert last nodes first
)
for i, (left, right) in enumerate(sorted_tuples):
if left:
import_node = _make_from_import_node(
left=left,
right=sorted(right),
trailing_nl=i == 0 and insert_pos == 0,
)
tree.insert_child(insert_pos, import_node)
else:
for j, name in enumerate(right):
import_node = _make_bare_import_node(
name=name,
trailing_nl=i == 0 and j == 0 and insert_pos == 0,
)
tree.insert_child(insert_pos, import_node)
def _create_simple_stmt_node_and_insert_behind(code, node):
if node is None or node.type != python_symbols.simple_stmt:
return
simple_stmt_node = Node(python_symbols.simple_stmt, [utils.newline_node(node)])
_node = utils.code_repr(code).children[0].children[0]
_node.parent = None
simple_stmt_node.insert_child(0, _node)
simple_stmt_node.prefix = utils.get_indent(node)
utils.insert_node_behind(node, simple_stmt_node)
def _remove_with_dygraph_guard(node: LN, capture: Capture,
filename: Filename):
# index of with_node, with_node will be replaced with simple_stmt node
with_node = capture['with']
parent = with_node.parent
idx = None
for i, child in enumerate(parent.children):
if child is with_node:
idx = i
break
# create simple_stmt node for "paddle.disable_static"
arg_list_nodes = capture['arg_list']
simple_stmt_disable_static = Node(python_symbols.simple_stmt,
[utils.newline_node(node)])
_node = utils.code_repr('paddle.disable_static' +
str(arg_list_nodes)).children[0].children[0]
_node.parent = None
simple_stmt_disable_static.insert_child(0, _node)
simple_stmt_disable_static.prefix = with_node.prefix
# create simple_stmt node for "paddle.enable_static"
simple_stmt_enable_static = Node(python_symbols.simple_stmt,
[utils.newline_node(node)])
simple_stmt_enable_static
_node = utils.code_repr(
'paddle.enable_static()').children[0].children[0]
_node.parent = None
simple_stmt_enable_static.insert_child(0, _node)
simple_stmt_enable_static.prefix = utils.get_indent(with_node)
suite_node = capture['suite']
# remove first newline
for node in suite_node.children:
if not isinstance(node, Leaf):
continue
if node.type == token.NEWLINE:
node.remove()
break
# remove first indent node, and add indent prefix to sibling node.
indent = None
for node in suite_node.children:
if not isinstance(node, Leaf):
continue
if node.type == token.INDENT:
indent = node.value
if node.next_sibling is not None:
node.next_sibling.prefix = node.prefix + indent
node.remove()
break
# transfer post leading dedent node prefix to sibling of with node
leaves = [leaf for leaf in suite_node.leaves()]
# visit all leaves in reversed order
last_dedent_leaf_idx = len(leaves)
for leaf in leaves[::-1]:
if leaf.type == token.DEDENT:
with_node.next_sibling.prefix = leaf.prefix + with_node.next_sibling.prefix
leaf.prefix = ""
else:
break
# remove dedenet node corresponding to with node
for node in suite_node.children[::-1]:
if not isinstance(node, Leaf):
continue
if node.type == token.DEDENT:
node.remove()
break
# unindent all code in suite
for node in suite_node.leaves():
if node.type == token.INDENT:
node.value = utils.dec_indent(node.value)
else:
node.prefix = utils.dec_indent(node.prefix)
with_node.remove()
parent.insert_child(idx, simple_stmt_disable_static)
idx += 1
for node in suite_node.children:
parent.insert_child(idx, node)
idx += 1
parent.insert_child(idx, simple_stmt_enable_static)
