Add Try node

Author: cdce8p

astroid/node_classes.py

@@ -70,6 +70,7 @@
     Slice,
     Starred,
     Subscript,
+    Try,
     TryExcept,
     TryFinally,
     Tuple,

astroid/nodes/__init__.py

@@ -92,6 +92,7 @@
     Starred,
     Statement,
     Subscript,
+    Try,
     TryExcept,
     TryFinally,
     Tuple,
@@ -204,6 +205,7 @@
     Slice,
     Starred,
     Subscript,
+    Try,
     TryExcept,
     TryFinally,
     Tuple,

astroid/nodes/node_classes.py

@@ -3953,6 +3953,129 @@ def get_children(self):
         yield self.slice
 
 
+class Try(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn, Statement):
+    """Class representing a :class:`ast.Try` node.
+
+    Replaces the `TryExcept` and `TryFinally` nodes.
+    Requires ``tree_rev >= 2``.
+
+    >>> import astroid
+    >>> node = astroid.extract_node('''
+        try:
+            do_something()
+        except Exception as error:
+            print("Error!")
+        ''')
+    >>> node
+    <Try l.2 at 0x7f23b2e41d68>
+    """
+
+    _astroid_fields = ("body", "handlers", "orelse", "finalbody")
+    _multi_line_block_fields = ("body", "handlers", "orelse", "finalbody")
+
+    def __init__(
+        self,
+        *,
+        lineno: Optional[int] = None,
+        col_offset: Optional[int] = None,
+        end_lineno: Optional[int] = None,
+        end_col_offset: Optional[int] = None,
+        parent: Optional[NodeNG] = None,
+    ) -> None:
+        """
+        :param lineno: The line that this node appears on in the source code.
+
+        :param col_offset: The column that this node appears on in the
+            source code.
+
+        :param parent: The parent node in the syntax tree.
+
+        :param end_lineno: The last line this node appears on in the source code.
+
+        :param end_col_offset: The end column this node appears on in the
+            source code. Note: This is after the last symbol.
+        """
+        self.body: typing.List[NodeNG] = []
+        """The contents of the block to catch exceptions from."""
+
+        self.handlers: typing.List[ExceptHandler] = []
+        """The exception handlers."""
+
+        self.orelse: typing.List[NodeNG] = []
+        """The contents of the ``else`` block."""
+
+        self.finalbody: typing.List[NodeNG] = []
+        """The contents of the ``finally`` block."""
+
+        super().__init__(
+            lineno=lineno,
+            col_offset=col_offset,
+            end_lineno=end_lineno,
+            end_col_offset=end_col_offset,
+            parent=parent,
+        )
+
+    def postinit(
+        self,
+        *,
+        body: Optional[typing.List[NodeNG]] = None,
+        handlers: Optional[typing.List[ExceptHandler]] = None,
+        orelse: Optional[typing.List[NodeNG]] = None,
+        finalbody: Optional[typing.List[NodeNG]] = None,
+    ) -> None:
+        """Do some setup after initialisation.
+
+        :param body: The contents of the block to catch exceptions from.
+
+        :param handlers: The exception handlers.
+
+        :param orelse: The contents of the ``else`` block.
+
+        :param finalbody: The contents of the ``finally`` block.
+        """
+        if body:
+            self.body = body
+        if handlers:
+            self.handlers = handlers
+        if orelse:
+            self.orelse = orelse
+        if finalbody:
+            self.finalbody = finalbody
+
+    def _infer_name(self, frame, name):
+        return name
+
+    def block_range(self, lineno: int) -> typing.Tuple[int, int]:
+        """Get a range from a given line number to where this node ends."""
+        if lineno == self.fromlineno:
+            return lineno, lineno
+        if self.body and self.body[0].fromlineno <= lineno <= self.body[-1].tolineno:
+            # Inside try body - return from lineno till end of try body
+            return lineno, self.body[-1].tolineno
+        for exhandler in self.handlers:
+            if exhandler.type and lineno == exhandler.type.fromlineno:
+                return lineno, lineno
+            if exhandler.body[0].fromlineno <= lineno <= exhandler.body[-1].tolineno:
+                return lineno, exhandler.body[-1].tolineno
+        if self.orelse:
+            if self.orelse[0].fromlineno - 1 == lineno:
+                return lineno, lineno
+            if self.orelse[0].fromlineno <= lineno <= self.orelse[-1].tolineno:
+                return lineno, self.orelse[-1].tolineno
+        if self.finalbody:
+            if self.finalbody[0].fromlineno - 1 == lineno:
+                return lineno, lineno
+            if self.finalbody[0].fromlineno <= lineno <= self.finalbody[-1].tolineno:
+                return lineno, self.finalbody[-1].tolineno
+        return lineno, self.tolineno
+
+    def get_children(self):
+        yield from self.body
+        yield from self.handlers
+        yield from self.orelse
+        yield from self.finalbody
+
+
 class TryExcept(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn, Statement):
     """Class representing an :class:`ast.TryExcept` node.
 

astroid/nodes/node_ng.py

@@ -587,7 +587,7 @@ def _get_yield_nodes_skip_lambdas(self):
         yield from ()
 
     def _infer_name(self, frame, name):
-        # overridden for ImportFrom, Import, Global, TryExcept and Arguments
+        # overridden for ImportFrom, Import, Global, Try, TryExcept and Arguments
         pass
 
     def _infer(self, context=None):

astroid/rebuilder.py

@@ -547,7 +547,7 @@ def visit(self, node: "ast.Starred", parent: NodeNG) -> nodes.Starred:
         @overload
         def visit(
             self, node: "ast.Try", parent: NodeNG
-        ) -> Union[nodes.TryExcept, nodes.TryFinally]:
+        ) -> Union[nodes.Try, nodes.TryExcept, nodes.TryFinally]:
             ...
 
         @overload
@@ -893,7 +893,7 @@ def visit(self, node: "ast.Starred", parent: NodeNG) -> nodes.Starred:
         @overload
         def visit(
             self, node: "ast.Try", parent: NodeNG
-        ) -> Union[nodes.TryExcept, nodes.TryFinally]:
+        ) -> Union[nodes.Try, nodes.TryExcept, nodes.TryFinally]:
             ...
 
         @overload
@@ -2134,28 +2134,45 @@ def visit_tryexcept(self, node: "ast.Try", parent: NodeNG) -> nodes.TryExcept:
 
     def visit_try(
         self, node: "ast.Try", parent: NodeNG
-    ) -> Union[nodes.TryExcept, nodes.TryFinally, None]:
-        # python 3.3 introduce a new Try node replacing
-        # TryFinally/TryExcept nodes
-        if node.finalbody:
-            newnode = nodes.TryFinally(
-                lineno=node.lineno,
-                col_offset=node.col_offset,
-                # end_lineno and end_col_offset added in 3.8
-                end_lineno=getattr(node, "end_lineno", None),
-                end_col_offset=getattr(node, "end_col_offset", None),
-                parent=parent,
-            )
-            body: List[Union[NodeNG, nodes.TryExcept]]
+    ) -> Union[nodes.Try, nodes.TryExcept, nodes.TryFinally, None]:
+        newnode: Union[nodes.Try, nodes.TryExcept, nodes.TryFinally]
+        if self._tree_rev == 1:
+            # python 3.3 introduce a new Try node replacing
+            # TryFinally/TryExcept nodes
+            if node.finalbody:
+                newnode = nodes.TryFinally(
+                    lineno=node.lineno,
+                    col_offset=node.col_offset,
+                    # end_lineno and end_col_offset added in 3.8
+                    end_lineno=getattr(node, "end_lineno", None),
+                    end_col_offset=getattr(node, "end_col_offset", None),
+                    parent=parent,
+                )
+                body: List[Union[NodeNG, nodes.TryExcept]]
+                if node.handlers:
+                    body = [self.visit_tryexcept(node, newnode)]
+                else:
+                    body = [self.visit(child, newnode) for child in node.body]
+                newnode.postinit(body, [self.visit(n, newnode) for n in node.finalbody])
+                return newnode
             if node.handlers:
-                body = [self.visit_tryexcept(node, newnode)]
-            else:
-                body = [self.visit(child, newnode) for child in node.body]
-            newnode.postinit(body, [self.visit(n, newnode) for n in node.finalbody])
-            return newnode
-        if node.handlers:
-            return self.visit_tryexcept(node, parent)
-        return None
+                return self.visit_tryexcept(node, parent)
+            return None
+
+        newnode = nodes.Try(
+            lineno=node.lineno,
+            col_offset=node.col_offset,
+            end_lineno=getattr(node, "end_lineno", None),
+            end_col_offset=getattr(node, "end_col_offset", None),
+            parent=parent,
+        )
+        newnode.postinit(
+            body=[self.visit(child, newnode) for child in node.body],
+            handlers=[self.visit(child, newnode) for child in node.handlers],
+            orelse=[self.visit(child, newnode) for child in node.orelse],
+            finalbody=[self.visit(child, newnode) for child in node.finalbody],
+        )
+        return newnode
 
     def visit_tuple(self, node: "ast.Tuple", parent: NodeNG) -> nodes.Tuple:
         """visit a Tuple node by returning a fresh instance of it"""

tests/unittest_nodes.py

@@ -61,6 +61,7 @@
     ParentMissingError,
     StatementMissing,
 )
+from astroid.manager import AstroidManager
 from astroid.nodes.node_classes import (
     AssignAttr,
     AssignName,
@@ -438,6 +439,38 @@ def test_if_typing_guard() -> None:
         assert code[3].is_typing_guard() is False
 
 
+class TryNodeTest(_NodeTest):
+    CODE = """
+        try:  # L2
+            print("Hello")
+        except IOError:
+            pass
+        except UnicodeError:
+            pass
+        else:
+            print()
+        finally:
+            print()
+    """
+
+    def test_block_range(self) -> None:
+        AstroidManager.tree_rev = 2
+        try:
+            assert self.astroid.body[0].block_range(1) == (1, 11)
+            assert self.astroid.body[0].block_range(2) == (2, 2)
+            assert self.astroid.body[0].block_range(3) == (3, 3)
+            assert self.astroid.body[0].block_range(4) == (4, 4)
+            assert self.astroid.body[0].block_range(5) == (5, 5)
+            assert self.astroid.body[0].block_range(6) == (6, 6)
+            assert self.astroid.body[0].block_range(7) == (7, 7)
+            assert self.astroid.body[0].block_range(8) == (8, 8)
+            assert self.astroid.body[0].block_range(9) == (9, 9)
+            assert self.astroid.body[0].block_range(10) == (10, 10)
+            assert self.astroid.body[0].block_range(11) == (11, 11)
+        finally:
+            AstroidManager.tree_rev = 1
+
+
 class TryExceptNodeTest(_NodeTest):
     CODE = """
         try: