Add MultiLineBlockMixin

`_multi_line_block_fields` is a list of strings indicating which
fields contain multi-line blocks. `_get_multi_line_blocks` dynamically
accesses these attributes the first time it is called and then caches
the resulting references so as to avoid repeated expensive attribute
lookups.

1 files changed, 35 insertions, 0 deletions

diff --git a/astroid/mixins.py b/astroid/mixins.py
index 739828d5..be5352f0 100644
--- a/astroid/mixins.py
+++ b/astroid/mixins.py
@@ -131,3 +131,38 @@ class ImportFromMixin(FilterStmtsMixin):
         raise exceptions.AttributeInferenceError(
             'Could not find original name for {attribute} in {target!r}',
             target=self, attribute=asname)
+
+
+class MultiLineBlockMixin:
+    """Mixin for nodes with multi-line blocks, e.g. For and FunctionDef.
+    Note that this does not apply to every node with a `body` field.
+    For instance, an If node has a multi-line body, but the body of an
+    IfExpr is not multi-line, and hence cannot contain Return nodes,
+    Assign nodes, etc.
+    """
+
+    @decorators.cachedproperty
+    def _multi_line_blocks(self):
+        return tuple(
+            getattr(self, field)
+            for field in self._multi_line_block_fields
+        )
+
+    def _get_return_nodes_skip_functions(self):
+        for block in self._multi_line_blocks:
+            for child_node in block:
+                if child_node.is_function:
+                    continue
+                yield from child_node._get_return_nodes_skip_functions()
+
+    def _get_yield_nodes_skip_lambdas(self):
+        for block in self._multi_line_blocks:
+            for child_node in block:
+                if child_node.is_lambda:
+                    continue
+                yield from child_node._get_yield_nodes_skip_lambdas()
+
+    def _get_assign_nodes(self):
+        for block in self._multi_line_blocks:
+            for child_node in block:
+                yield from child_node._get_assign_nodes()