|
6 | 6 | order to get a single Astroid representation
|
7 | 7 | """
|
8 | 8 |
|
9 |
| -import ast |
10 | 9 | import sys
|
11 | 10 | import token
|
12 | 11 | import tokenize
|
13 | 12 | from io import StringIO
|
14 | 13 | from tokenize import TokenInfo, generate_tokens
|
15 | 14 | from typing import (
|
| 15 | + TYPE_CHECKING, |
16 | 16 | Callable,
|
17 | 17 | Dict,
|
18 | 18 | Generator,
|
|
39 | 39 | else:
|
40 | 40 | from typing_extensions import Final
|
41 | 41 |
|
| 42 | +if TYPE_CHECKING: |
| 43 | + import ast |
| 44 | + |
42 | 45 |
|
43 | 46 | REDIRECT: Final[Dict[str, str]] = {
|
44 | 47 | "arguments": "Arguments",
|
@@ -1386,18 +1389,25 @@ def _find_orelse_keyword(
|
1386 | 1389 | if not self._data or not node.orelse:
|
1387 | 1390 | return None, None
|
1388 | 1391 |
|
1389 |
| - # If the first child in orelse is an If node the orelse is an elif block |
1390 |
| - if isinstance(node.orelse[0], ast.If): |
1391 |
| - return node.orelse[0].lineno, node.orelse[0].col_offset |
1392 |
| - |
1393 |
| - end_lineno = node.orelse[0].lineno - 1 |
| 1392 | + end_lineno = node.orelse[0].lineno |
| 1393 | + |
| 1394 | + def find_keyword(begin: int, end: int) -> Tuple[Optional[int], Optional[int]]: |
| 1395 | + # pylint: disable-next=unsubscriptable-object |
| 1396 | + data = "\n".join(self._data[begin:end]) |
| 1397 | + |
| 1398 | + try: |
| 1399 | + tokens = list(generate_tokens(StringIO(data).readline)) |
| 1400 | + except tokenize.TokenError: |
| 1401 | + # If we cut-off in the middle of multi-line if statements we |
| 1402 | + # generate a TokenError here. We just keep trying |
| 1403 | + # until the multi-line statement is closed. |
| 1404 | + return find_keyword(begin, end + 1) |
| 1405 | + for t in tokens[::-1]: |
| 1406 | + if t.type == token.NAME and t.string in {"else", "elif"}: |
| 1407 | + return node.lineno + t.start[0] - 1, t.start[1] |
| 1408 | + return None, None |
1394 | 1409 |
|
1395 |
| - # pylint: disable-next=unsubscriptable-object |
1396 |
| - data = "\n".join(self._data[node.lineno - 1 : end_lineno]) |
1397 |
| - for t in generate_tokens(StringIO(data).readline): |
1398 |
| - if t.type == token.NAME and t.string == "else": |
1399 |
| - return node.lineno + t.start[0] - 1, t.start[1] |
1400 |
| - return None, None |
| 1410 | + return find_keyword(node.lineno - 1, end_lineno) |
1401 | 1411 |
|
1402 | 1412 | def visit_if(self, node: "ast.If", parent: NodeNG) -> nodes.If:
|
1403 | 1413 | """visit an If node by returning a fresh instance of it"""
|
|
0 commit comments