Improve _get_return_nodes_skip_functions algorithm

Return statements can only appear in multiline bodies (like function
definitions), so there is no need to check other nodes for them.

Author: nickdrozd

astroid/node_classes.py

@@ -651,11 +651,7 @@ def _get_name_nodes(self):
                 yield matching
 
     def _get_return_nodes_skip_functions(self):
-        for child_node in self.get_children():
-            if child_node.is_function:
-                continue
-            for matching in child_node._get_return_nodes_skip_functions():
-                yield matching
+        yield from ()
 
     def _get_yield_nodes_skip_lambdas(self):
         for child_node in self.get_children():
@@ -2832,6 +2828,12 @@ def catch(self, exceptions): # pylint: disable=redefined-outer-name
                 return True
         return False
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
 
 class Exec(Statement):
     """Class representing the ``exec`` statement.
@@ -2982,6 +2984,17 @@ def _get_assign_nodes(self):
         for child_node in self.orelse:
             yield from child_node._get_assign_nodes()
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
+        for child_node in self.orelse:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
 
 class AsyncFor(For):
     """Class representing an :class:`ast.AsyncFor` node.
@@ -3262,6 +3275,17 @@ def _get_assign_nodes(self):
         for child_node in self.orelse:
             yield from child_node._get_assign_nodes()
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
+        for child_node in self.orelse:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
 
 class IfExp(NodeNG):
     """Class representing an :class:`ast.IfExp` node.
@@ -3673,12 +3697,6 @@ def get_children(self):
     def _get_return_nodes_skip_functions(self):
         yield self
 
-        for child_node in self.get_children():
-            if child_node.is_function:
-                continue
-            for matching in child_node._get_return_nodes_skip_functions():
-                yield matching
-
 
 class Set(_BaseContainer):
     """Class representing an :class:`ast.Set` node.
@@ -3987,6 +4005,18 @@ def _get_assign_nodes(self):
         for child_node in self.orelse:
             yield from child_node._get_assign_nodes()
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
+        for child_node in self.orelse or ():
+            if child_node.is_function:
+                continue
+            for matching in child_node._get_return_nodes_skip_functions():
+                yield matching
+
 
 class TryFinally(mixins.BlockRangeMixIn, Statement):
     """Class representing an :class:`ast.TryFinally` node.
@@ -4054,6 +4084,18 @@ def _get_assign_nodes(self):
         for child_node in self.finalbody:
             yield from child_node._get_assign_nodes()
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
+        for child_node in self.finalbody:
+            if child_node.is_function:
+                continue
+            for matching in child_node._get_return_nodes_skip_functions():
+                yield matching
+
 
 class Tuple(_BaseContainer):
     """Class representing an :class:`ast.Tuple` node.
@@ -4252,6 +4294,18 @@ def _get_assign_nodes(self):
         for child_node in self.orelse:
             yield from child_node._get_assign_nodes()
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
+        for child_node in self.orelse:
+            if child_node.is_function:
+                continue
+            for matching in child_node._get_return_nodes_skip_functions():
+                yield matching
+
 
 class With(mixins.BlockRangeMixIn, mixins.AssignTypeMixin, Statement):
     """Class representing an :class:`ast.With` node.
@@ -4313,6 +4367,12 @@ def _get_assign_nodes(self):
         for child_node in self.body:
             yield from child_node._get_assign_nodes()
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
 
 class AsyncWith(With):
     """Asynchronous ``with`` built with the ``async`` keyword."""

astroid/scoped_nodes.py

@@ -1598,6 +1598,12 @@ def _get_assign_nodes(self):
         for child_node in self.body:
             yield from child_node._get_assign_nodes()
 
+    def _get_return_nodes_skip_functions(self):
+        for child_node in self.body:
+            if child_node.is_function:
+                continue
+            yield from child_node._get_return_nodes_skip_functions()
+
 
 class AsyncFunctionDef(FunctionDef):
     """Class representing an :class:`ast.FunctionDef` node.