summaryrefslogtreecommitdiff
path: root/astroid/mixins.py
diff options
context:
space:
mode:
Diffstat (limited to 'astroid/mixins.py')
-rw-r--r--astroid/mixins.py35
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()