@@ -734,9 +734,42 @@ def on_conditional_branch():
734734 self .report (messages .UndefinedName , node , name )
735735
736736 def handleChildren (self , tree , omit = None ):
737+ """Handle all children recursively, but may be flattened."""
737738 for node in iter_child_nodes (tree , omit = omit ):
738739 self .handleNode (node , tree )
739740
741+ def handleChildrenNested (self , node ):
742+ """Handle all children recursively."""
743+ self .handleChildren (node )
744+
745+ def _iter_flattened (self , tree , omit , _fields_order = _FieldsOrder ()):
746+ """
747+ Yield child nodes of *node* and their children, with handler.
748+
749+ The value yielded is a tuple of the node, its parent and its handler.
750+ The handler may be False to indicate that no handler and no recursion
751+ is required as the node is part of a flattened list.
752+ """
753+ _may_flatten = (self .handleChildren ,
754+ self .handleChildrenFlattened )
755+
756+ nodes = [(tree , None )]
757+ for node , parent in nodes :
758+ # Skip the root of the tree, which has parent None
759+ handler = self .getNodeHandler (node .__class__ ) if parent else False
760+ if handler and handler not in _may_flatten :
761+ yield node , parent , handler
762+ else :
763+ nodes [:] += ((child , node )
764+ for child in iter_child_nodes (node ,
765+ omit ,
766+ _fields_order ))
767+
768+ def handleChildrenFlattened (self , tree , omit = None ):
769+ """Handle all children recursively as a flat list where possible."""
770+ for node , parent , handler in self ._iter_flattened (tree , omit = omit ):
771+ self .handleNode (node , parent , handler )
772+
740773 def isLiteralTupleUnpacking (self , node ):
741774 if isinstance (node , ast .Assign ):
742775 for child in node .targets + [node .value ]:
@@ -766,7 +799,12 @@ def getDocstring(self, node):
766799
767800 return (node .s , doctest_lineno )
768801
769- def handleNode (self , node , parent ):
802+ def handleNode (self , node , parent , handler = None ):
803+ """
804+ Handle a single node, invoking its handler, which may recurse.
805+
806+ If handler is None, the default handler is used.
807+ """
770808 if node is None :
771809 return
772810 if self .offset and getattr (node , 'lineno' , None ) is not None :
@@ -777,11 +815,18 @@ def handleNode(self, node, parent):
777815 if self .futuresAllowed and not (isinstance (node , ast .ImportFrom ) or
778816 self .isDocstring (node )):
779817 self .futuresAllowed = False
780- self . nodeDepth += 1
781- node .depth = self .nodeDepth
818+
819+ node .depth = self .nodeDepth + 1
782820 node .parent = parent
783- try :
821+
822+ if handler is False :
823+ return
824+
825+ if not handler :
784826 handler = self .getNodeHandler (node .__class__ )
827+
828+ self .nodeDepth += 1
829+ try :
785830 handler (node )
786831 finally :
787832 self .nodeDepth -= 1
@@ -833,21 +878,22 @@ def ignore(self, node):
833878 pass
834879
835880 # "stmt" type nodes
836- DELETE = PRINT = FOR = ASYNCFOR = WHILE = IF = WITH = WITHITEM = \
837- ASYNCWITH = ASYNCWITHITEM = RAISE = TRYFINALLY = EXEC = \
838- EXPR = ASSIGN = handleChildren
881+ DELETE = PRINT = EXEC = EXPR = RAISE = handleChildrenFlattened
882+ ASSIGN = TRYFINALLY = handleChildren
883+ FOR = ASYNCFOR = WHILE = IF = WITH = ASYNCWITH = handleChildren
884+ WITHITEM = ASYNCWITHITEM = handleChildrenFlattened
839885
840886 PASS = ignore
841887
842888 # "expr" type nodes
843889 BOOLOP = BINOP = UNARYOP = IFEXP = DICT = SET = \
844890 COMPARE = CALL = REPR = ATTRIBUTE = SUBSCRIPT = \
845- STARRED = NAMECONSTANT = handleChildren
891+ STARRED = NAMECONSTANT = handleChildrenFlattened
846892
847893 NUM = STR = BYTES = ELLIPSIS = ignore
848894
849895 # "slice" type nodes
850- SLICE = EXTSLICE = INDEX = handleChildren
896+ SLICE = EXTSLICE = INDEX = handleChildrenFlattened
851897
852898 # expression contexts are node instances too, though being constants
853899 LOAD = STORE = DEL = AUGLOAD = AUGSTORE = PARAM = ignore
@@ -859,7 +905,8 @@ def ignore(self, node):
859905 MATMULT = ignore
860906
861907 # additional node types
862- COMPREHENSION = KEYWORD = FORMATTEDVALUE = handleChildren
908+ COMPREHENSION = handleChildren
909+ KEYWORD = FORMATTEDVALUE = handleChildrenFlattened
863910
864911 def ASSERT (self , node ):
865912 if isinstance (node .test , ast .Tuple ) and node .test .elts != []:
@@ -903,7 +950,7 @@ def GENERATOREXP(self, node):
903950 self .handleChildren (node )
904951 self .popScope ()
905952
906- LISTCOMP = handleChildren if PY2 else GENERATOREXP
953+ LISTCOMP = handleChildrenNested if PY2 else GENERATOREXP
907954
908955 DICTCOMP = SETCOMP = GENERATOREXP
909956
0 commit comments