Return a mapping from {lineno: "assertion test expression"}.
(src: bytes)
| 550 | |
| 551 | @functools.lru_cache(maxsize=1) |
| 552 | def _get_assertion_exprs(src: bytes) -> dict[int, str]: |
| 553 | """Return a mapping from {lineno: "assertion test expression"}.""" |
| 554 | ret: dict[int, str] = {} |
| 555 | |
| 556 | depth = 0 |
| 557 | lines: list[str] = [] |
| 558 | assert_lineno: int | None = None |
| 559 | seen_lines: set[int] = set() |
| 560 | |
| 561 | def _write_and_reset() -> None: |
| 562 | nonlocal depth, lines, assert_lineno, seen_lines |
| 563 | assert assert_lineno is not None |
| 564 | ret[assert_lineno] = "".join(lines).rstrip().rstrip("\\") |
| 565 | depth = 0 |
| 566 | lines = [] |
| 567 | assert_lineno = None |
| 568 | seen_lines = set() |
| 569 | |
| 570 | tokens = tokenize.tokenize(io.BytesIO(src).readline) |
| 571 | for tp, source, (lineno, offset), _, line in tokens: |
| 572 | if tp == tokenize.NAME and source == "assert": |
| 573 | assert_lineno = lineno |
| 574 | elif assert_lineno is not None: |
| 575 | # keep track of depth for the assert-message `,` lookup |
| 576 | if tp == tokenize.OP and source in "([{": |
| 577 | depth += 1 |
| 578 | elif tp == tokenize.OP and source in ")]}": |
| 579 | depth -= 1 |
| 580 | |
| 581 | if not lines: |
| 582 | lines.append(line[offset:]) |
| 583 | seen_lines.add(lineno) |
| 584 | # a non-nested comma separates the expression from the message |
| 585 | elif depth == 0 and tp == tokenize.OP and source == ",": |
| 586 | # one line assert with message |
| 587 | if lineno in seen_lines and len(lines) == 1: |
| 588 | offset_in_trimmed = offset + len(lines[-1]) - len(line) |
| 589 | lines[-1] = lines[-1][:offset_in_trimmed] |
| 590 | # multi-line assert with message |
| 591 | elif lineno in seen_lines: |
| 592 | lines[-1] = lines[-1][:offset] |
| 593 | # multi line assert with escaped newline before message |
| 594 | else: |
| 595 | lines.append(line[:offset]) |
| 596 | _write_and_reset() |
| 597 | elif tp in {tokenize.NEWLINE, tokenize.ENDMARKER}: |
| 598 | _write_and_reset() |
| 599 | elif lines and lineno not in seen_lines: |
| 600 | lines.append(line) |
| 601 | seen_lines.add(lineno) |
| 602 | |
| 603 | return ret |
| 604 | |
| 605 | |
| 606 | class AssertionRewriter(ast.NodeVisitor): |