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.

2 files changed, 77 insertions, 11 deletions

diff --git a/astroid/node_classes.py b/astroid/node_classes.py
index 1297ba29..e802b7fc 100644
--- a/astroid/node_classes.py
+++ b/astroid/node_classes.py
@@ -651,11 +651,7 @@ class NodeNG(object):
                 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 @@ class ExceptHandler(mixins.AssignTypeMixin, Statement):
                 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 @@ class For(mixins.BlockRangeMixIn, mixins.AssignTypeMixin, Statement):
         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 @@ class If(mixins.BlockRangeMixIn, Statement):
         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 @@ class Return(Statement):
     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 @@ class TryExcept(mixins.BlockRangeMixIn, Statement):
         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 @@ class TryFinally(mixins.BlockRangeMixIn, Statement):
         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 @@ class While(mixins.BlockRangeMixIn, Statement):
         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 @@ class With(mixins.BlockRangeMixIn, mixins.AssignTypeMixin, Statement):
         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."""
diff --git a/astroid/scoped_nodes.py b/astroid/scoped_nodes.py
index 66a036bb..1e491abd 100644
--- a/astroid/scoped_nodes.py
+++ b/astroid/scoped_nodes.py
@@ -1598,6 +1598,12 @@ class FunctionDef(node_classes.Statement, Lambda):
         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.