Add more tests and include more nodes

Author: DanielNoord

ChangeLog

@@ -8,7 +8,8 @@ Release date: TBA
 
 * ``astroid`` now requires Python 3.7.2 to run.
 
-* Add ``orelse_lineno`` and ``orelse_col_offset`` attributes to ``nodes.If``.
+* Add ``orelse_lineno`` and ``orelse_col_offset`` attributes to ``nodes.If``,
+  ``nodes.For``, ``nodes.While``, and ``nodes.TryExcept``.
 
 * Fix ``re`` brain on Python ``3.11``. The flags now come from ``re._compile``.
 

astroid/nodes/node_classes.py

@@ -2661,6 +2661,12 @@ def __init__(
         self.orelse: typing.List[NodeNG] = []
         """The contents of the ``else`` block of the loop."""
 
+        self.orelse_lineno: Optional[int] = None
+        """The line number of the ``else`` keyword."""
+
+        self.orelse_col_offset: Optional[int] = None
+        """The column offset of the ``else`` keyword."""
+
         self.type_annotation: Optional[NodeNG] = None  # can be None
         """If present, this will contain the type annotation passed by a type comment"""
 
@@ -2680,6 +2686,9 @@ def postinit(
         body: Optional[typing.List[NodeNG]] = None,
         orelse: Optional[typing.List[NodeNG]] = None,
         type_annotation: Optional[NodeNG] = None,
+        *,
+        orelse_lineno: Optional[int] = None,
+        orelse_col_offset: Optional[int] = None,
     ) -> None:
         """Do some setup after initialisation.
 
@@ -2690,13 +2699,19 @@ def postinit(
         :param body: The contents of the body of the loop.
 
         :param orelse: The contents of the ``else`` block of the loop.
+
+        :param orelse_lineno: The line number of the ``else`` keyword.
+
+        :param orelse_lineno: The column offset of the ``else`` keyword.
         """
         self.target = target
         self.iter = iter
         if body is not None:
             self.body = body
         if orelse is not None:
             self.orelse = orelse
+        self.orelse_lineno = orelse_lineno
+        self.orelse_col_offset = orelse_col_offset
         self.type_annotation = type_annotation
 
     assigned_stmts: ClassVar[AssignedStmtsCall["For"]]
@@ -3979,6 +3994,12 @@ def __init__(
         self.orelse: typing.List[NodeNG] = []
         """The contents of the ``else`` block."""
 
+        self.orelse_lineno: Optional[int] = None
+        """The line number of the ``else`` or ``elif`` keyword."""
+
+        self.orelse_col_offset: Optional[int] = None
+        """The column offset of the ``else`` or ``elif`` keyword."""
+
         super().__init__(
             lineno=lineno,
             col_offset=col_offset,
@@ -3992,6 +4013,9 @@ def postinit(
         body: Optional[typing.List[NodeNG]] = None,
         handlers: Optional[typing.List[ExceptHandler]] = None,
         orelse: Optional[typing.List[NodeNG]] = None,
+        *,
+        orelse_lineno: Optional[int] = None,
+        orelse_col_offset: Optional[int] = None,
     ) -> None:
         """Do some setup after initialisation.
 
@@ -4000,13 +4024,19 @@ def postinit(
         :param handlers: The exception handlers.
 
         :param orelse: The contents of the ``else`` block.
+
+        :param orelse_lineno: The line number of the ``else`` keyword.
+
+        :param orelse_lineno: The column offset of the ``else`` keyword.
         """
         if body is not None:
             self.body = body
         if handlers is not None:
             self.handlers = handlers
         if orelse is not None:
             self.orelse = orelse
+        self.orelse_lineno = orelse_lineno
+        self.orelse_col_offset = orelse_col_offset
 
     def _infer_name(self, frame, name):
         return name
@@ -4341,6 +4371,12 @@ def __init__(
         self.orelse: typing.List[NodeNG] = []
         """The contents of the ``else`` block."""
 
+        self.orelse_lineno: Optional[int] = None
+        """The line number of the ``else`` or ``elif`` keyword."""
+
+        self.orelse_col_offset: Optional[int] = None
+        """The column offset of the ``else`` or ``elif`` keyword."""
+
         super().__init__(
             lineno=lineno,
             col_offset=col_offset,
@@ -4354,6 +4390,9 @@ def postinit(
         test: Optional[NodeNG] = None,
         body: Optional[typing.List[NodeNG]] = None,
         orelse: Optional[typing.List[NodeNG]] = None,
+        *,
+        orelse_lineno: Optional[int] = None,
+        orelse_col_offset: Optional[int] = None,
     ) -> None:
         """Do some setup after initialisation.
 
@@ -4362,12 +4401,18 @@ def postinit(
         :param body: The contents of the loop.
 
         :param orelse: The contents of the ``else`` block.
+
+        :param orelse_lineno: The line number of the ``else`` keyword.
+
+        :param orelse_lineno: The column offset of the ``else`` keyword.
         """
         self.test = test
         if body is not None:
             self.body = body
         if orelse is not None:
             self.orelse = orelse
+        self.orelse_lineno = orelse_lineno
+        self.orelse_col_offset = orelse_col_offset
 
     @cached_property
     def blockstart_tolineno(self):

astroid/rebuilder.py

@@ -248,6 +248,26 @@ def _reset_end_lineno(self, newnode: nodes.NodeNG) -> None:
         for child_node in newnode.get_children():
             self._reset_end_lineno(child_node)
 
+    def _find_orelse_keyword(
+        self, node: Union["ast.If", "ast.Try", "ast.For", "ast.AsyncFor", "ast.While"]
+    ) -> Tuple[Optional[int], Optional[int]]:
+        """Get the line number and column offset of the `else` or `elif` keyword."""
+        if not self._data or not node.orelse:
+            return None, None
+
+        start = node.orelse[0].lineno
+
+        # pylint: disable-next=unsubscriptable-object
+        for index, line in enumerate(self._data[start - 1 :: -1]):
+            if line.lstrip().startswith("else"):
+                # if start - index == 37:
+                # breakpoint()
+                return start - index, line.index("else")
+            if line.lstrip().startswith("elif"):
+                return start - index, line.index("elif")
+
+        return None, None
+
     def visit_module(
         self, node: "ast.Module", modname: str, modpath: str, package: bool
     ) -> nodes.Module:
@@ -1175,6 +1195,8 @@ def _visit_for(
             # pylint: disable-next=unsubscriptable-object
             col_offset = self._data[node.lineno - 1].index("async")
 
+        orelse_lineno, orelse_col_offset = self._find_orelse_keyword(node)
+
         newnode = cls(
             lineno=node.lineno,
             col_offset=col_offset,
@@ -1190,6 +1212,8 @@ def _visit_for(
             body=[self.visit(child, newnode) for child in node.body],
             orelse=[self.visit(child, newnode) for child in node.orelse],
             type_annotation=type_annotation,
+            orelse_lineno=orelse_lineno,
+            orelse_col_offset=orelse_col_offset,
         )
         return newnode
 
@@ -1371,24 +1395,6 @@ def visit_global(self, node: "ast.Global", parent: NodeNG) -> nodes.Global:
                 self._global_names[-1].setdefault(name, []).append(newnode)
         return newnode
 
-    def _find_orelse_keyword(
-        self, node: "ast.If"
-    ) -> Tuple[Optional[int], Optional[int]]:
-        """Get the line number and column offset of the `else` or `elif` keyword."""
-        if not self._data or not node.orelse:
-            return None, None
-
-        start = node.orelse[0].lineno
-
-        # pylint: disable-next=unsubscriptable-object
-        for index, line in enumerate(self._data[start:node.lineno:-1]):
-            if line.rstrip().startswith("else"):
-                return start - index + 1, line.index("else")
-            if line.rstrip().startswith("elif"):
-                return start - index + 1, line.index("elif")
-
-        return None, None
-
     def visit_if(self, node: "ast.If", parent: NodeNG) -> nodes.If:
         """visit an If node by returning a fresh instance of it"""
         newnode = nodes.If(
@@ -1820,10 +1826,15 @@ def visit_tryexcept(self, node: "ast.Try", parent: NodeNG) -> nodes.TryExcept:
             )
         else:
             newnode = nodes.TryExcept(node.lineno, node.col_offset, parent)
+
+        orelse_lineno, orelse_col_offset = self._find_orelse_keyword(node)
+
         newnode.postinit(
             [self.visit(child, newnode) for child in node.body],
             [self.visit(child, newnode) for child in node.handlers],
             [self.visit(child, newnode) for child in node.orelse],
+            orelse_lineno=orelse_lineno,
+            orelse_col_offset=orelse_col_offset,
         )
         return newnode
 
@@ -1891,10 +1902,15 @@ def visit_while(self, node: "ast.While", parent: NodeNG) -> nodes.While:
             end_col_offset=getattr(node, "end_col_offset", None),
             parent=parent,
         )
+
+        orelse_lineno, orelse_col_offset = self._find_orelse_keyword(node)
+
         newnode.postinit(
             self.visit(node.test, newnode),
             [self.visit(child, newnode) for child in node.body],
             [self.visit(child, newnode) for child in node.orelse],
+            orelse_lineno=orelse_lineno,
+            orelse_col_offset=orelse_col_offset,
         )
         return newnode
 

tests/unittest_nodes.py

@@ -334,12 +334,52 @@ class IfNodeTest(_NodeTest):
             # This is using else in a comment
             raise
 
+        for x in range(3):
+            print()
+        else:
+            print()
+
+        for x in range(3):
+            print()
+            if x == 3:
+                break
+            else:
+                print()
+        else:
+            print()
+
+        while True:
+            print()
+        else:
+            print()
+
+        try:
+            1 / 0
+        except ZeroDivisionError:
+            print()
+        else:
+            print()
+        finally:
+            print()
+
+        try:
+            1 / 0
+        except ZeroDivisionError:
+            try:
+                1 / 0
+            except:
+                print()
+            else:
+                print()
+        else:
+            print()
+
     """
 
     def test_if_elif_else_node(self) -> None:
         """test transformation for If node"""
-        self.assertEqual(len(self.astroid.body), 5)
-        for stmt in self.astroid.body:
+        self.assertEqual(len(self.astroid.body), 10)
+        for stmt in self.astroid.body[:5]:
             self.assertIsInstance(stmt, nodes.If)
         self.assertFalse(self.astroid.body[0].orelse)  # simple If
         self.assertIsInstance(self.astroid.body[1].orelse[0], nodes.Pass)  # If / else
@@ -348,8 +388,8 @@ def test_if_elif_else_node(self) -> None:
 
     def test_block_range(self) -> None:
         # XXX ensure expected values
-        self.assertEqual(self.astroid.block_range(1), (0, 33))
-        self.assertEqual(self.astroid.block_range(10), (0, 33))  # XXX (10, 33) ?
+        self.assertEqual(self.astroid.block_range(1), (0, 73))
+        self.assertEqual(self.astroid.block_range(10), (0, 73))  # XXX (10, 73) ?
         self.assertEqual(self.astroid.body[1].block_range(5), (5, 6))
         self.assertEqual(self.astroid.body[1].block_range(6), (6, 6))
         self.assertEqual(self.astroid.body[1].orelse[0].block_range(7), (7, 8))
@@ -373,6 +413,20 @@ def test_orelse_line_numbering(self) -> None:
         assert self.astroid.body[4].orelse_col_offset == 0
         assert self.astroid.body[4].orelse[0].orelse_lineno == 31
         assert self.astroid.body[4].orelse[0].orelse_col_offset == 0
+        assert self.astroid.body[5].orelse_lineno == 37
+        assert self.astroid.body[5].orelse_col_offset == 0
+        assert self.astroid.body[6].orelse_lineno == 46
+        assert self.astroid.body[6].orelse_col_offset == 0
+        assert self.astroid.body[6].body[1].orelse_lineno == 44
+        assert self.astroid.body[6].body[1].orelse_col_offset == 4
+        assert self.astroid.body[7].orelse_lineno == 51
+        assert self.astroid.body[7].orelse_col_offset == 0
+        assert self.astroid.body[8].body[0].orelse_lineno == 58
+        assert self.astroid.body[8].body[0].orelse_col_offset == 0
+        assert self.astroid.body[9].orelse_lineno == 72
+        assert self.astroid.body[9].orelse_col_offset == 0
+        assert self.astroid.body[9].handlers[0].body[0].orelse_lineno == 70
+        assert self.astroid.body[9].handlers[0].body[0].orelse_col_offset == 4
 
     @staticmethod
     @pytest.mark.filterwarnings("ignore:.*is_sys_guard:DeprecationWarning")