# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# [Link]
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
======================================================================
========
if isinstance(node, list):
return [[Link](n) for n in node]
elif isinstance(node, tuple):
return tuple([Link](n) for n in node)
elif not isinstance(node, ([Link], [Link])):
# Assuming everything that's not an AST, list or tuple is a value type
# and may simply be assigned.
return node
"""AST manipulation utilities."""
import ast
import gast
from [Link] import anno
from [Link] import parser
from [Link] import qual_names
class CleanCopier(object):
"""NodeTransformer-like visitor that copies an AST."""
def __init__(self, preserve_annos):
super(CleanCopier, self).__init__()
self.preserve_annos = preserve_annos
def copy(self, node):
"""Returns a deep copy of node (excluding some fields, see copy_clean)."""
if isinstance(node, list):
return [[Link](n) for n in node]
elif isinstance(node, tuple):
return tuple([Link](n) for n in node)
elif not isinstance(node, ([Link], [Link])):
# Assuming everything that's not an AST, list or tuple is a value type
# and may simply be assigned.
return node
assert isinstance(node, ([Link], [Link]))
new_fields = {}
for f in node._fields:
if not [Link]('__') and hasattr(node, f):
new_fields[f] = [Link](getattr(node, f))
new_node = type(node)(**new_fields)
if self.preserve_annos:
for k in self.preserve_annos:
[Link](node, new_node, k)
return new_node
def copy_clean(node, preserve_annos=None):
"""Creates a deep copy of an AST.
The copy will not include fields that are prefixed by '__', with the
exception of user-specified annotations.
Args:
node: [Link]
preserve_annos: Optional[Set[Hashable]], annotation keys to include in the
copy
Returns:
[Link]
"""
return CleanCopier(preserve_annos).copy(node)
class SymbolRenamer([Link]):
"""Transformer that can rename symbols to a simple names."""
def __init__(self, name_map):
self.name_map = name_map
def _process_name_node(self, node):
qn = [Link](node, [Link])
if qn in self.name_map:
new_node = [Link](
str(self.name_map[qn]),
ctx=[Link],
annotation=None,
type_comment=None)
# All annotations get carried over.
for k in [Link](node):
[Link](node, new_node, k)
return new_node
return self.generic_visit(node)
def _process_list_of_strings(self, names):
for i in range(len(names)):
qn = qual_names.QN(names[i])
if qn in self.name_map:
names[i] = str(self.name_map[qn])
return names
def visit_Nonlocal(self, node):
[Link] = self._process_list_of_strings([Link])
return node
def visit_Global(self, node):
[Link] = self._process_list_of_strings([Link])
return node
def visit_Name(self, node):
return self._process_name_node(node)
def visit_Attribute(self, node):
if [Link](node, [Link]):
return self._process_name_node(node)
# Renaming attributes is not supported.
return self.generic_visit(node)
def visit_FunctionDef(self, node):
qn = qual_names.QN([Link])
if qn in self.name_map:
[Link] = str(self.name_map[qn])
return self.generic_visit(node)
def rename_symbols(node, name_map):
"""Renames symbols in an AST. Requires qual_names annotations."""
renamer = SymbolRenamer(name_map)
if isinstance(node, list):
return [[Link](n) for n in node]
elif isinstance(node, tuple):
return tuple([Link](n) for n in node)
return [Link](node)
def keywords_to_dict(keywords):
"""Converts a list of [Link] objects to a dict."""
keys = []
values = []
for kw in keywords:
[Link]([Link]([Link], kind=None))
[Link]([Link])
return [Link](keys=keys, values=values)
class PatternMatcher([Link]):
"""Matches a node against a pattern represented by a node."""
def __init__(self, pattern):
[Link] = pattern
self.pattern_stack = []
[Link] = True
def compare_and_visit(self, node, pattern):
self.pattern_stack.append([Link])
[Link] = pattern
self.generic_visit(node)
[Link] = self.pattern_stack.pop()
def no_match(self):
[Link] = False
return False
def is_wildcard(self, p):
if isinstance(p, (list, tuple)) and len(p) == 1:
p, = p
if isinstance(p, [Link]) and [Link] == '_':
return True
if p == '_':
return True
return False
def generic_visit(self, node):
if not [Link]:
return
pattern = [Link]
for f in node._fields:
if [Link]('__'):
continue
if not hasattr(node, f):
if hasattr(pattern, f) and getattr(pattern, f):
return self.no_match()
else:
continue
if not hasattr(pattern, f):
return self.no_match()
v = getattr(node, f)
p = getattr(pattern, f)
if self.is_wildcard(p):
continue
if isinstance(v, (list, tuple)):
if not isinstance(p, (list, tuple)) or len(v) != len(p):
return self.no_match()
for v_item, p_item in zip(v, p):
self.compare_and_visit(v_item, p_item)
elif isinstance(v, ([Link], [Link])):
if not isinstance(v, type(p)) and not isinstance(p, type(v)):
return self.no_match()
self.compare_and_visit(v, p)
else:
# Assume everything else is a value type.
if v != p:
return self.no_match()
def matches(node, pattern):
"""Basic pattern matcher for AST.
The pattern may contain wildcards represented by the symbol '_'. A node
matches a pattern if for every node in the tree, either there is a node of
the same type in pattern, or a Name node with id='_'.
Args:
node: [Link]
pattern: [Link]
Returns:
bool
"""
if isinstance(pattern, str):
pattern = parser.parse_str(pattern)
matcher = PatternMatcher(pattern)
[Link](node)
return [Link]
# TODO(mdan): Once we have error tracing, we may be able to just go to SSA.
def apply_to_single_assignments(targets, values, apply_fn):
"""Applies a function to each individual assignment.
This function can process a possibly-unpacked (e.g. a, b = c, d) assignment.
It tries to break down the unpacking if possible. In effect, it has the same
effect as passing the assigned values in SSA form to apply_fn.
Examples:
The following will result in apply_fn(a, c), apply_fn(b, d):
a, b = c, d
The following will result in apply_fn(a, c[0]), apply_fn(b, c[1]):
a, b = c
The following will result in apply_fn(a, (b, c)):
a = b, c
It uses the visitor pattern to allow subclasses to process single
assignments individually.
Args:
targets: Union[List[[Link], ...], Tuple[[Link], ...], [Link], should be
used with the targets field of an [Link] node
values: [Link]
apply_fn: Callable[[[Link], [Link]], None], called with the
respective nodes of each single assignment
"""
if not isinstance(targets, (list, tuple)):
targets = (targets,)
for target in targets:
if isinstance(target, ([Link], [Link])):
for i in range(len([Link])):
target_el = [Link][i]
if isinstance(values, ([Link], [Link])):
value_el = [Link][i]
else:
idx = parser.parse_expression(str(i))
value_el = [Link](values, idx, ctx=[Link]())
apply_to_single_assignments(target_el, value_el, apply_fn)
else:
apply_fn(target, values)
def parallel_walk(node, other):
"""Walks two ASTs in parallel.
The two trees must have identical structure.
Args:
node: Union[[Link], Iterable[[Link]]]
other: Union[[Link], Iterable[[Link]]]
Yields:
Tuple[[Link], [Link]]
Raises:
ValueError: if the two trees don't have identical structure.
"""
if isinstance(node, (list, tuple)):
node_stack = list(node)
else:
node_stack = [node]
if isinstance(other, (list, tuple)):
other_stack = list(other)
else:
other_stack = [other]
while node_stack and other_stack:
assert len(node_stack) == len(other_stack)
n = node_stack.pop()
o = other_stack.pop()
if ((not isinstance(n, ([Link], [Link], str)) and n is not None) or
(not isinstance(o, ([Link], [Link], str)) and n is not None) or
n.__class__.__name__ != o.__class__.__name__):
raise ValueError('inconsistent nodes: {} ({}) and {} ({})'.format(
n, n.__class__.__name__, o, o.__class__.__name__))
yield n, o
if isinstance(n, str):
assert isinstance(o, str), 'The check above should have ensured this'
continue
if n is None:
assert o is None, 'The check above should have ensured this'
continue
for f in n._fields:
n_child = getattr(n, f, None)
o_child = getattr(o, f, None)
if [Link]('__') or n_child is None or o_child is None:
continue
if isinstance(n_child, (list, tuple)):
if (not isinstance(o_child, (list, tuple)) or
len(n_child) != len(o_child)):
raise ValueError(
'inconsistent values for field {}: {} and {}'.format(
f, n_child, o_child))
node_stack.extend(n_child)
other_stack.extend(o_child)
elif isinstance(n_child, ([Link], [Link])):
node_stack.append(n_child)
other_stack.append(o_child)
elif n_child != o_child:
raise ValueError(
'inconsistent values for field {}: {} and {}'.format(
f, n_child, o_child))