Improve _get_yield_nodes_skip_lambdas algorithm

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

Author: nickdrozd

astroid/node_classes.py

@@ -654,11 +654,7 @@ def _get_return_nodes_skip_functions(self):
         yield from ()
 
     def _get_yield_nodes_skip_lambdas(self):
-        for child_node in self.get_children():
-            if child_node.is_lambda:
-                continue
-            for matching in child_node._get_yield_nodes_skip_lambdas():
-                yield matching
+        yield from ()
 
     def _infer_name(self, frame, name):
         # overridden for ImportFrom, Import, Global, TryExcept and Arguments
@@ -2708,6 +2704,10 @@ def postinit(self, value=None):
     def get_children(self):
         yield self.value
 
+    def _get_yield_nodes_skip_lambdas(self):
+        if not self.value.is_lambda:
+            yield from self.value._get_yield_nodes_skip_lambdas()
+
 
 class Ellipsis(NodeNG): # pylint: disable=redefined-builtin
     """Class representing an :class:`ast.Ellipsis` node.
@@ -2834,6 +2834,12 @@ def _get_return_nodes_skip_functions(self):
                 continue
             yield from child_node._get_return_nodes_skip_functions()
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class Exec(Statement):
     """Class representing the ``exec`` statement.
@@ -2995,6 +3001,17 @@ def _get_return_nodes_skip_functions(self):
                 continue
             yield from child_node._get_return_nodes_skip_functions()
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
+        for child_node in self.orelse:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class AsyncFor(For):
     """Class representing an :class:`ast.AsyncFor` node.
@@ -3286,6 +3303,17 @@ def _get_return_nodes_skip_functions(self):
                 continue
             yield from child_node._get_return_nodes_skip_functions()
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
+        for child_node in self.orelse:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class IfExp(NodeNG):
     """Class representing an :class:`ast.IfExp` node.
@@ -4017,6 +4045,17 @@ def _get_return_nodes_skip_functions(self):
             for matching in child_node._get_return_nodes_skip_functions():
                 yield matching
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
+        for child_node in self.orelse:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class TryFinally(mixins.BlockRangeMixIn, Statement):
     """Class representing an :class:`ast.TryFinally` node.
@@ -4096,6 +4135,17 @@ def _get_return_nodes_skip_functions(self):
             for matching in child_node._get_return_nodes_skip_functions():
                 yield matching
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
+        for child_node in self.finalbody:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class Tuple(_BaseContainer):
     """Class representing an :class:`ast.Tuple` node.
@@ -4306,6 +4356,17 @@ def _get_return_nodes_skip_functions(self):
             for matching in child_node._get_return_nodes_skip_functions():
                 yield matching
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
+        for child_node in self.orelse:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class With(mixins.BlockRangeMixIn, mixins.AssignTypeMixin, Statement):
     """Class representing an :class:`ast.With` node.
@@ -4373,6 +4434,12 @@ def _get_return_nodes_skip_functions(self):
                 continue
             yield from child_node._get_return_nodes_skip_functions()
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class AsyncWith(With):
     """Asynchronous ``with`` built with the ``async`` keyword."""
@@ -4407,12 +4474,6 @@ def get_children(self):
     def _get_yield_nodes_skip_lambdas(self):
         yield self
 
-        for child_node in self.get_children():
-            if child_node.is_function_or_lambda:
-                continue
-            for matching in child_node._get_yield_nodes_skip_lambdas():
-                yield matching
-
 
 class YieldFrom(Yield):
     """Class representing an :class:`ast.YieldFrom` node."""

astroid/scoped_nodes.py

@@ -1604,6 +1604,12 @@ def _get_return_nodes_skip_functions(self):
                 continue
             yield from child_node._get_return_nodes_skip_functions()
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 class AsyncFunctionDef(FunctionDef):
     """Class representing an :class:`ast.FunctionDef` node.
@@ -2714,6 +2720,12 @@ def _get_assign_nodes(self):
         for child_node in self.body:
             yield from child_node._get_assign_nodes()
 
+    def _get_yield_nodes_skip_lambdas(self):
+        for child_node in self.body:
+            if child_node.is_lambda:
+                continue
+            yield from child_node._get_yield_nodes_skip_lambdas()
+
 
 # Backwards-compatibility aliases
 Class = util.proxy_alias('Class', ClassDef)