| 143 | """Ensure AST positions are preserved during rewriting (#12818).""" |
| 144 | |
| 145 | def preserved(code: str) -> None: |
| 146 | s = textwrap.dedent(code) |
| 147 | locations = [] |
| 148 | |
| 149 | def loc(msg: str | None = None) -> None: |
| 150 | frame = inspect.currentframe() |
| 151 | assert frame |
| 152 | frame = frame.f_back |
| 153 | assert frame |
| 154 | frame = frame.f_back |
| 155 | assert frame |
| 156 | |
| 157 | offset = frame.f_lasti |
| 158 | |
| 159 | instructions = {i.offset: i for i in dis.get_instructions(frame.f_code)} |
| 160 | |
| 161 | # skip CACHE instructions |
| 162 | while offset not in instructions and offset >= 0: |
| 163 | offset -= 1 |
| 164 | |
| 165 | instruction = instructions[offset] |
| 166 | if sys.version_info >= (3, 11): |
| 167 | position = instruction.positions |
| 168 | else: |
| 169 | position = instruction.starts_line |
| 170 | |
| 171 | locations.append((msg, instruction.opname, position)) |
| 172 | |
| 173 | globals = {"loc": loc} |
| 174 | |
| 175 | m = rewrite(s) |
| 176 | mod = compile(m, "<string>", "exec") |
| 177 | exec(mod, globals, globals) |
| 178 | transformed_locations = locations |
| 179 | locations = [] |
| 180 | |
| 181 | mod = compile(s, "<string>", "exec") |
| 182 | exec(mod, globals, globals) |
| 183 | original_locations = locations |
| 184 | |
| 185 | assert len(original_locations) > 0 |
| 186 | assert original_locations == transformed_locations |
| 187 | |
| 188 | preserved(""" |
| 189 | def f(): |