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