Making Pylint faster 2

astroid/mixins.py

@@ -131,3 +131,38 @@ def real_name(self, asname):
         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()

astroid/node_classes.py

@@ -643,28 +643,18 @@ def nodes_of_class(self, klass, skip_klass=None):
                 yield matching
 
     def _get_assign_nodes(self):
-        for child_node in self.get_children():
-            for matching in child_node._get_assign_nodes():
-                yield matching
+        yield from ()
 
     def _get_name_nodes(self):
         for child_node in self.get_children():
             for matching in child_node._get_name_nodes():
                 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():
-            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
@@ -1737,9 +1727,7 @@ def get_children(self):
     def _get_assign_nodes(self):
         yield self
 
-        for child_node in self.get_children():
-            for matching in child_node._get_assign_nodes():
-                yield matching
+        yield from self.value._get_assign_nodes()
 
 
 class AnnAssign(mixins.AssignTypeMixin, Statement):
@@ -2716,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.
@@ -2749,7 +2741,8 @@ def get_children(self):
         yield from ()
 
 
-class ExceptHandler(mixins.AssignTypeMixin, Statement):
+class ExceptHandler(mixins.MultiLineBlockMixin,
+                    mixins.AssignTypeMixin, Statement):
     """Class representing an :class:`ast.ExceptHandler`. node.
 
     An :class:`ExceptHandler` is an ``except`` block on a try-except.
@@ -2766,6 +2759,7 @@ class ExceptHandler(mixins.AssignTypeMixin, Statement):
     [<ExceptHandler l.4 at 0x7f23b2e9e860>]
     """
     _astroid_fields = ('type', 'name', 'body',)
+    _multi_line_block_fields = ('body',)
     type = None
     """The types that the block handles.
 
@@ -2906,14 +2900,16 @@ def postinit(self, dims=None):
         self.dims = dims
 
 
-class For(mixins.BlockRangeMixIn, mixins.AssignTypeMixin, Statement):
+class For(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn,
+          mixins.AssignTypeMixin, Statement):
     """Class representing an :class:`ast.For` node.
 
     >>> node = astroid.extract_node('for thing in things: print(thing)')
     >>> node
     <For l.1 at 0x7f23b2e8cf28>
     """
     _astroid_fields = ('target', 'iter', 'body', 'orelse',)
+    _multi_line_block_fields = ('body', 'orelse')
     target = None
     """What the loop assigns to.
 
@@ -3181,14 +3177,15 @@ def get_children(self):
         yield from ()
 
 
-class If(mixins.BlockRangeMixIn, Statement):
+class If(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn, Statement):
     """Class representing an :class:`ast.If` node.
 
     >>> node = astroid.extract_node('if condition: print(True)')
     >>> node
     <If l.1 at 0x7f23b2e9dd30>
     """
     _astroid_fields = ('test', 'body', 'orelse')
+    _multi_line_block_fields = ('body', 'orelse')
     test = None
     """The condition that the statement tests.
 
@@ -3663,12 +3660,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.
@@ -3896,7 +3887,7 @@ def get_children(self):
         yield self.slice
 
 
-class TryExcept(mixins.BlockRangeMixIn, Statement):
+class TryExcept(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn, Statement):
     """Class representing an :class:`ast.TryExcept` node.
 
     >>> node = astroid.extract_node('''
@@ -3909,6 +3900,7 @@ class TryExcept(mixins.BlockRangeMixIn, Statement):
     <TryExcept l.2 at 0x7f23b2e9d908>
     """
     _astroid_fields = ('body', 'handlers', 'orelse',)
+    _multi_line_block_fields = ('body', 'orelse')
     body = None
     """The contents of the block to catch exceptions from.
 
@@ -3971,7 +3963,8 @@ def get_children(self):
         yield from self.orelse or ()
 
 
-class TryFinally(mixins.BlockRangeMixIn, Statement):
+class TryFinally(mixins.MultiLineBlockMixin,
+                 mixins.BlockRangeMixIn, Statement):
     """Class representing an :class:`ast.TryFinally` node.
 
     >>> node = astroid.extract_node('''
@@ -3986,6 +3979,7 @@ class TryFinally(mixins.BlockRangeMixIn, Statement):
     <TryFinally l.2 at 0x7f23b2e41d68>
     """
     _astroid_fields = ('body', 'finalbody',)
+    _multi_line_block_fields = ('body', 'finalbody')
     body = None
     """The try-except that the finally is attached to.
 
@@ -4152,7 +4146,7 @@ def get_children(self):
         yield self.operand
 
 
-class While(mixins.BlockRangeMixIn, Statement):
+class While(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn, Statement):
     """Class representing an :class:`ast.While` node.
 
     >>> node = astroid.extract_node('''
@@ -4163,6 +4157,7 @@ class While(mixins.BlockRangeMixIn, Statement):
     <While l.2 at 0x7f23b2e4e390>
     """
     _astroid_fields = ('test', 'body', 'orelse',)
+    _multi_line_block_fields = ('body', 'orelse')
     test = None
     """The condition that the loop tests.
 
@@ -4222,7 +4217,8 @@ def get_children(self):
         yield from self.orelse
 
 
-class With(mixins.BlockRangeMixIn, mixins.AssignTypeMixin, Statement):
+class With(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn,
+           mixins.AssignTypeMixin, Statement):
     """Class representing an :class:`ast.With` node.
 
     >>> node = astroid.extract_node('''
@@ -4233,6 +4229,7 @@ class With(mixins.BlockRangeMixIn, mixins.AssignTypeMixin, Statement):
     <With l.2 at 0x7f23b2e4e710>
     """
     _astroid_fields = ('items', 'body')
+    _multi_line_block_fields = ('body',)
     items = None
     """The pairs of context managers and the names they are assigned to.
 
@@ -4312,12 +4309,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

@@ -1205,7 +1205,7 @@ def get_children(self):
         yield self.body
 
 
-class FunctionDef(node_classes.Statement, Lambda):
+class FunctionDef(mixins.MultiLineBlockMixin, node_classes.Statement, Lambda):
     """Class representing an :class:`ast.FunctionDef`.
 
     >>> node = astroid.extract_node('''
@@ -1216,6 +1216,7 @@ class FunctionDef(node_classes.Statement, Lambda):
     <FunctionDef.my_func l.2 at 0x7f23b2e71e10>
     """
     _astroid_fields = ('decorators', 'args', 'returns', 'body')
+    _multi_line_block_fields = ('body',)
     returns = None
     decorators = None
     """The decorators that are applied to this method or function.
@@ -1712,7 +1713,6 @@ def get_wrapping_class(node):
     return klass
 
 
-
 class ClassDef(mixins.FilterStmtsMixin, LocalsDictNodeNG,
                node_classes.Statement):
     """Class representing an :class:`ast.ClassDef` node.
@@ -2701,6 +2701,10 @@ def get_children(self):
         if self.decorators is not None:
             yield self.decorators
 
+    def _get_assign_nodes(self):
+        for child_node in self.body:
+            yield from child_node._get_assign_nodes()
+
 
 # Backwards-compatibility aliases
 Class = util.proxy_alias('Class', ClassDef)