@@ -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