def test_parallel_walk_inconsistent_trees(self): node_1 = parser.parse_str( textwrap.dedent(""" def f(a): return a + 1 """)) node_2 = parser.parse_str( textwrap.dedent(""" def f(a): return a + (a * 2) """)) node_3 = parser.parse_str( textwrap.dedent(""" def f(a): return a + 2 """)) with self.assertRaises(ValueError): for _ in ast_util.parallel_walk(node_1, node_2): pass # There is not particular reason to reject trees that differ only in the # value of a constant. # TODO(mdan): This should probably be allowed. with self.assertRaises(ValueError): for _ in ast_util.parallel_walk(node_1, node_3): pass
def test_parallel_walk(self): node = parser.parse_str( textwrap.dedent(""" def f(a): return a + 1 """)) for child_a, child_b in ast_util.parallel_walk(node, node): self.assertEqual(child_a, child_b)
def test_parallel_walk_string_leaves(self): src = """ def f(a): global g """ node = parser.parse(textwrap.dedent(src)) for child_a, child_b in ast_util.parallel_walk(node, node): self.assertEqual(child_a, child_b)
def create_source_map(nodes, code, filename, indices_in_code): """Creates a source map between an annotated AST and the code it compiles to. Args: nodes: Iterable[ast.AST, ...] code: Text filename: Optional[Text] indices_in_code: Union[int, Iterable[int, ...]], the positions at which nodes appear in code. The parser always returns a module when parsing code. This argument indicates the position in that module's body at which the corresponding of node should appear. Returns: Dict[CodeLocation, OriginInfo], mapping locations in code to locations indicated by origin annotations in node. """ reparsed_nodes = parser.parse_str(code) reparsed_nodes = [reparsed_nodes.body[i] for i in indices_in_code] resolve(reparsed_nodes, code) result = {} for before, after in ast_util.parallel_walk(nodes, reparsed_nodes): # Note: generated code might not be mapped back to its origin. # TODO(mdan): Generated code should always be mapped to something. origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None) final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None) if origin_info is None or final_info is None: continue line_loc = LineLocation(filename, final_info.loc.lineno) existing_origin = result.get(line_loc) if existing_origin is not None: # Overlaps may exist because of child nodes, but almost never to # different line locations. Exception make decorated functions, where # both lines are mapped to the same line in the AST. # Line overlaps: keep bottom node. if existing_origin.loc.line_loc == origin_info.loc.line_loc: if existing_origin.loc.lineno >= origin_info.loc.lineno: continue # In case of overlaps, keep the leftmost node. if existing_origin.loc.col_offset <= origin_info.loc.col_offset: continue result[line_loc] = origin_info return result
def create_source_map(nodes, code, filename, indices_in_code): """Creates a source map between an annotated AST and the code it compiles to. Args: nodes: Iterable[ast.AST, ...] code: Text filename: Optional[Text] indices_in_code: Union[int, Iterable[int, ...]], the positions at which nodes appear in code. The parser always returns a module when parsing code. This argument indicates the position in that module's body at which the corresponding of node should appear. Returns: Dict[LineLocation, OriginInfo], mapping locations in code to locations indicated by origin annotations in node. """ reparsed_nodes = parser.parse_str(code) reparsed_nodes = [reparsed_nodes.body[i] for i in indices_in_code] for node in reparsed_nodes: resolve(node, code) result = {} try: for before, after in ast_util.parallel_walk(nodes, reparsed_nodes): # Note: generated code might not be mapped back to its origin. # TODO(mdan): Generated code should always be mapped to something. origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None) final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None) if origin_info is None or final_info is None: continue line_loc = LineLocation(filename, final_info.loc.lineno) existing_origin = result.get(line_loc) if existing_origin is not None: # Overlaps may exist because of child nodes, but almost never to # different line locations. Exception make decorated functions, where # both lines are mapped to the same line in the AST. # Line overlaps: keep bottom node. if existing_origin.loc.line_loc == origin_info.loc.line_loc: if existing_origin.loc.lineno >= origin_info.loc.lineno: continue # In case of overlaps, keep the leftmost node. if existing_origin.loc.col_offset <= origin_info.loc.col_offset: continue result[line_loc] = origin_info except ValueError: if logging.has_verbosity(3): for n, rn in zip(nodes, reparsed_nodes): nodes_str = pretty_printer.fmt(n, color=False, noanno=True) reparsed_nodes_str = pretty_printer.fmt(rn, color=False, noanno=True) diff = difflib.context_diff(nodes_str.split('\n'), reparsed_nodes_str.split('\n'), fromfile='Original nodes', tofile='Reparsed nodes', n=7) diff = '\n'.join(diff) logging.log(3, 'AST seems to lack integrity. Diff:\n%s', diff) raise return result
def create_source_map(nodes, code, filepath): """Creates a source map between an annotated AST and the code it compiles to. Note: this function assumes nodes nodes, code and filepath correspond to the same code. Args: nodes: Iterable[ast.AST, ...], one or more AST modes. code: Text, the source code in which nodes are found. filepath: Text Returns: Dict[LineLocation, OriginInfo], mapping locations in code to locations indicated by origin annotations in node. """ reparsed_nodes = parser.parse_str(code, preamble_len=0, single_node=False) for node in reparsed_nodes: resolve(node, code, filepath, node.lineno, node.col_offset) source_map = {} try: for before, after in ast_util.parallel_walk(nodes, reparsed_nodes): # Note: generated code might not be mapped back to its origin. # TODO(mdan): Generated code should always be mapped to something. origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None) final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None) if origin_info is None or final_info is None: continue # Note: the keys are by line only, excluding the column offset. line_loc = LineLocation(final_info.loc.filename, final_info.loc.lineno) existing_origin = source_map.get(line_loc) if existing_origin is not None: # Overlaps may exist because of child nodes, but almost never to # different line locations. Exception make decorated functions, where # both lines are mapped to the same line in the AST. # Line overlaps: keep bottom node. if existing_origin.loc.line_loc == origin_info.loc.line_loc: if existing_origin.loc.lineno >= origin_info.loc.lineno: continue # In case of column overlaps, keep the leftmost node. if existing_origin.loc.col_offset <= origin_info.loc.col_offset: continue source_map[line_loc] = origin_info except ValueError: if logging.has_verbosity(3): for n, rn in zip(nodes, reparsed_nodes): nodes_str = pretty_printer.fmt(n, color=False, noanno=True) reparsed_nodes_str = pretty_printer.fmt(rn, color=False, noanno=True) diff = difflib.context_diff( nodes_str.split('\n'), reparsed_nodes_str.split('\n'), fromfile='Original nodes', tofile='Reparsed nodes', n=7) diff = '\n'.join(diff) logging.log(3, 'AST seems to lack integrity. Diff:\n%s', diff) raise return source_map
def create_source_map(nodes, code, filename): """Creates a source map between an annotated AST and the code it compiles to. Args: nodes: Iterable[ast.AST, ...] code: Text filename: Optional[Text] Returns: Dict[LineLocation, OriginInfo], mapping locations in code to locations indicated by origin annotations in node. """ reparsed_nodes = parser.parse_str(code, preamble_len=0, single_node=False) for node in reparsed_nodes: resolve(node, code) result = {} try: for before, after in ast_util.parallel_walk(nodes, reparsed_nodes): # Note: generated code might not be mapped back to its origin. # TODO(mdan): Generated code should always be mapped to something. origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None) final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None) if origin_info is None or final_info is None: continue line_loc = LineLocation(filename, final_info.loc.lineno) existing_origin = result.get(line_loc) if existing_origin is not None: # Overlaps may exist because of child nodes, but almost never to # different line locations. Exception make decorated functions, where # both lines are mapped to the same line in the AST. # Line overlaps: keep bottom node. if existing_origin.loc.line_loc == origin_info.loc.line_loc: if existing_origin.loc.lineno >= origin_info.loc.lineno: continue # In case of overlaps, keep the leftmost node. if existing_origin.loc.col_offset <= origin_info.loc.col_offset: continue result[line_loc] = origin_info except ValueError: if logging.has_verbosity(3): for n, rn in zip(nodes, reparsed_nodes): nodes_str = pretty_printer.fmt(n, color=False, noanno=True) reparsed_nodes_str = pretty_printer.fmt(rn, color=False, noanno=True) diff = difflib.context_diff( nodes_str.split('\n'), reparsed_nodes_str.split('\n'), fromfile='Original nodes', tofile='Reparsed nodes', n=7) diff = '\n'.join(diff) logging.log(3, 'AST seems to lack integrity. Diff:\n%s', diff) raise return result