Add type-specific nodes_of_class

nodes_of_class is a very flexible method, which is great for use in
client code (e.g. Pylint). However, that flexibility requires a great
deal of runtime type checking:

    def nodes_of_class(self, klass, skip_klass=None):
        if isinstance(self, klass):
            yield self

        if skip_klass is None:
            for child_node in self.get_children():
                for matching in child_node.nodes_of_class(klass, skip_klass):
                    yield matching

            return

        for child_node in self.get_children():
            if isinstance(child_node, skip_klass):
                continue
            for matching in child_node.nodes_of_class(klass, skip_klass):
                yield matching

First, the node has to check its own type to see whether it's of the
desired class. Then the skip_klass flag has to be checked to see
whether anything needs to be skipped. If so, the type of every yielded
node has to be check to see if it should be skipped.

This is fine for calling code whose arguments can't be known in
advance ("Give me all the Assign and ClassDef nodes, but skip all the
BinOps, YieldFroms, and Globals."), but in Astroid itself, every call
to this function can be known in advance. There's no need to do any
type checking if all the nodes know how to respond to certain
requests. Take get_assign_nodes for example. The Assign nodes know
that they should yield themselves and then yield their Assign
children. Other nodes know in advance that they aren't Assign nodes,
so they don't need to check their own type, just immediately yield
their Assign children.

Overly specific functions like get_yield_nodes_skip_lambdas certainly
aren't very elegant, but the tradeoff is to take advantage of knowing
how the library code works to improve speed.

Author: nickdrozd

astroid/node_classes.py

@@ -220,6 +220,7 @@ class NodeNG(object):
 
     :type: bool
     """
+    is_lambda = False
     # Attributes below are set by the builder module or by raw factories
     lineno = None
     """The line that this node appears on in the source code.
@@ -641,6 +642,30 @@ def nodes_of_class(self, klass, skip_klass=None):
             for matching in child_node.nodes_of_class(klass, skip_klass):
                 yield matching
 
+    def _get_assign_nodes(self):
+        for child_node in self.get_children():
+            for matching in child_node._get_assign_nodes():
+                yield matching
+
+    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
+
+    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
+
     def _infer_name(self, frame, name):
         # overridden for ImportFrom, Import, Global, TryExcept and Arguments
         return None
@@ -1267,6 +1292,13 @@ def __init__(self, name=None, lineno=None, col_offset=None, parent=None):
     def get_children(self):
         yield from ()
 
+    def _get_name_nodes(self):
+        yield self
+
+        for child_node in self.get_children():
+            for matching in child_node._get_name_nodes():
+                yield matching
+
 
 class Arguments(mixins.AssignTypeMixin, NodeNG):
     """Class representing an :class:`ast.arguments` node.
@@ -1702,6 +1734,13 @@ def get_children(self):
 
         yield self.value
 
+    def _get_assign_nodes(self):
+        yield self
+
+        for child_node in self.get_children():
+            for matching in child_node._get_assign_nodes():
+                yield matching
+
 
 class AnnAssign(mixins.AssignTypeMixin, Statement):
     """Class representing an :class:`ast.AnnAssign` node.
@@ -2792,7 +2831,7 @@ def catch(self, exceptions): # pylint: disable=redefined-outer-name
         """
         if self.type is None or exceptions is None:
             return True
-        for node in self.type.nodes_of_class(Name):
+        for node in self.type._get_name_nodes():
             if node.name in exceptions:
                 return True
         return False
@@ -3582,7 +3621,7 @@ def raises_not_implemented(self):
         """
         if not self.exc:
             return False
-        for name in self.exc.nodes_of_class(Name):
+        for name in self.exc._get_name_nodes():
             if name.name == 'NotImplementedError':
                 return True
         return False
@@ -3621,6 +3660,15 @@ def get_children(self):
         if self.value is not None:
             yield self.value
 
+    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.
@@ -4261,6 +4309,15 @@ def get_children(self):
         if self.value is not None:
             yield self.value
 
+    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

@@ -1037,6 +1037,7 @@ class Lambda(mixins.FilterStmtsMixin, LocalsDictNodeNG):
     _astroid_fields = ('args', 'body',)
     _other_other_fields = ('locals',)
     name = '<lambda>'
+    is_lambda = True
 
     # function's type, 'function' | 'method' | 'staticmethod' | 'classmethod'
     @property
@@ -1313,7 +1314,7 @@ def extra_decorators(self):
             return []
 
         decorators = []
-        for assign in frame.nodes_of_class(node_classes.Assign):
+        for assign in frame._get_assign_nodes():
             if (isinstance(assign.value, node_classes.Call)
                     and isinstance(assign.value.func, node_classes.Name)):
                 for assign_node in assign.targets:
@@ -1530,9 +1531,7 @@ def is_generator(self):
         :returns: True is this is a generator function, False otherwise.
         :rtype: bool
         """
-        yield_nodes = (node_classes.Yield, node_classes.YieldFrom)
-        return next(self.nodes_of_class(yield_nodes,
-                                        skip_klass=(FunctionDef, Lambda)), False)
+        return next(self._get_yield_nodes_skip_lambdas(), False)
 
     def infer_call_result(self, caller=None, context=None):
         """Infer what the function returns when called.
@@ -1563,7 +1562,7 @@ def infer_call_result(self, caller=None, context=None):
                 c._metaclass = metaclass
                 yield c
                 return
-        returns = self.nodes_of_class(node_classes.Return, skip_klass=FunctionDef)
+        returns = self._get_return_nodes_skip_functions()
         for returnnode in returns:
             if returnnode.value is None:
                 yield node_classes.Const(None)