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. fil
(nodes, code, filepath)
| 88 | |
| 89 | # TODO(mdan): This source map should be a class - easier to refer to. |
| 90 | def create_source_map(nodes, code, filepath): |
| 91 | """Creates a source map between an annotated AST and the code it compiles to. |
| 92 | |
| 93 | Note: this function assumes nodes nodes, code and filepath correspond to the |
| 94 | same code. |
| 95 | |
| 96 | Args: |
| 97 | nodes: Iterable[ast.AST, ...], one or more AST modes. |
| 98 | code: Text, the source code in which nodes are found. |
| 99 | filepath: Text |
| 100 | |
| 101 | Returns: |
| 102 | Dict[LineLocation, OriginInfo], mapping locations in code to locations |
| 103 | indicated by origin annotations in node. |
| 104 | """ |
| 105 | reparsed_nodes = parser.parse_str(code, preamble_len=0, single_node=False) |
| 106 | for node in reparsed_nodes: |
| 107 | resolve(node, code, filepath, node.lineno, node.col_offset) |
| 108 | |
| 109 | source_map = {} |
| 110 | |
| 111 | try: |
| 112 | for before, after in ast_util.parallel_walk(nodes, reparsed_nodes): |
| 113 | # Note: generated code might not be mapped back to its origin. |
| 114 | # TODO(mdan): Generated code should always be mapped to something. |
| 115 | origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None) |
| 116 | final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None) |
| 117 | if origin_info is None or final_info is None: |
| 118 | continue |
| 119 | |
| 120 | # Note: the keys are by line only, excluding the column offset. |
| 121 | line_loc = LineLocation(final_info.loc.filename, final_info.loc.lineno) |
| 122 | |
| 123 | existing_origin = source_map.get(line_loc) |
| 124 | if existing_origin is not None: |
| 125 | # Overlaps may exist because of child nodes, but almost never to |
| 126 | # different line locations. Exception make decorated functions, where |
| 127 | # both lines are mapped to the same line in the AST. |
| 128 | |
| 129 | # Line overlaps: keep bottom node. |
| 130 | if existing_origin.loc.line_loc == origin_info.loc.line_loc: |
| 131 | if existing_origin.loc.lineno >= origin_info.loc.lineno: |
| 132 | continue |
| 133 | |
| 134 | # In case of column overlaps, keep the leftmost node. |
| 135 | if existing_origin.loc.col_offset <= origin_info.loc.col_offset: |
| 136 | continue |
| 137 | |
| 138 | source_map[line_loc] = origin_info |
| 139 | |
| 140 | except ValueError: |
| 141 | if logging.has_verbosity(3): |
| 142 | for n, rn in zip(nodes, reparsed_nodes): |
| 143 | nodes_str = pretty_printer.fmt(n, color=False, noanno=True) |
| 144 | reparsed_nodes_str = pretty_printer.fmt(rn, color=False, noanno=True) |
| 145 | diff = difflib.context_diff( |
| 146 | nodes_str.split('\n'), |
| 147 | reparsed_nodes_str.split('\n'), |