@@ -734,9 +734,42 @@ def on_conditional_branch():
734
734
self .report (messages .UndefinedName , node , name )
735
735
736
736
def handleChildren (self , tree , omit = None ):
737
+ """Handle all children recursively, but may be flattened."""
737
738
for node in iter_child_nodes (tree , omit = omit ):
738
739
self .handleNode (node , tree )
739
740
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
+
740
773
def isLiteralTupleUnpacking (self , node ):
741
774
if isinstance (node , ast .Assign ):
742
775
for child in node .targets + [node .value ]:
@@ -766,7 +799,12 @@ def getDocstring(self, node):
766
799
767
800
return (node .s , doctest_lineno )
768
801
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
+ """
770
808
if node is None :
771
809
return
772
810
if self .offset and getattr (node , 'lineno' , None ) is not None :
@@ -777,11 +815,18 @@ def handleNode(self, node, parent):
777
815
if self .futuresAllowed and not (isinstance (node , ast .ImportFrom ) or
778
816
self .isDocstring (node )):
779
817
self .futuresAllowed = False
780
- self . nodeDepth += 1
781
- node .depth = self .nodeDepth
818
+
819
+ node .depth = self .nodeDepth + 1
782
820
node .parent = parent
783
- try :
821
+
822
+ if handler is False :
823
+ return
824
+
825
+ if not handler :
784
826
handler = self .getNodeHandler (node .__class__ )
827
+
828
+ self .nodeDepth += 1
829
+ try :
785
830
handler (node )
786
831
finally :
787
832
self .nodeDepth -= 1
@@ -833,21 +878,22 @@ def ignore(self, node):
833
878
pass
834
879
835
880
# "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
839
885
840
886
PASS = ignore
841
887
842
888
# "expr" type nodes
843
889
BOOLOP = BINOP = UNARYOP = IFEXP = DICT = SET = \
844
890
COMPARE = CALL = REPR = ATTRIBUTE = SUBSCRIPT = \
845
- STARRED = NAMECONSTANT = handleChildren
891
+ STARRED = NAMECONSTANT = handleChildrenFlattened
846
892
847
893
NUM = STR = BYTES = ELLIPSIS = ignore
848
894
849
895
# "slice" type nodes
850
- SLICE = EXTSLICE = INDEX = handleChildren
896
+ SLICE = EXTSLICE = INDEX = handleChildrenFlattened
851
897
852
898
# expression contexts are node instances too, though being constants
853
899
LOAD = STORE = DEL = AUGLOAD = AUGSTORE = PARAM = ignore
@@ -859,7 +905,8 @@ def ignore(self, node):
859
905
MATMULT = ignore
860
906
861
907
# additional node types
862
- COMPREHENSION = KEYWORD = FORMATTEDVALUE = handleChildren
908
+ COMPREHENSION = handleChildren
909
+ KEYWORD = FORMATTEDVALUE = handleChildrenFlattened
863
910
864
911
def ASSERT (self , node ):
865
912
if isinstance (node .test , ast .Tuple ) and node .test .elts != []:
@@ -903,7 +950,7 @@ def GENERATOREXP(self, node):
903
950
self .handleChildren (node )
904
951
self .popScope ()
905
952
906
- LISTCOMP = handleChildren if PY2 else GENERATOREXP
953
+ LISTCOMP = handleChildrenNested if PY2 else GENERATOREXP
907
954
908
955
DICTCOMP = SETCOMP = GENERATOREXP
909
956
0 commit comments