@@ -820,9 +820,42 @@ def on_conditional_branch():
820
820
self .report (messages .UndefinedName , node , name )
821
821
822
822
def handleChildren (self , tree , omit = None ):
823
+ """Handle all children recursively, but may be flattened."""
823
824
for node in iter_child_nodes (tree , omit = omit ):
824
825
self .handleNode (node , tree )
825
826
827
+ def handleChildrenNested (self , node ):
828
+ """Handle all children recursively."""
829
+ self .handleChildren (node )
830
+
831
+ def _iter_flattened (self , tree , omit , _fields_order = _FieldsOrder ()):
832
+ """
833
+ Yield child nodes of *node* and their children, with handler.
834
+
835
+ The value yielded is a tuple of the node, its parent and its handler.
836
+ The handler may be False to indicate that no handler and no recursion
837
+ is required as the node is part of a flattened list.
838
+ """
839
+ _may_flatten = (self .handleChildren ,
840
+ self .handleChildrenFlattened )
841
+
842
+ nodes = [(tree , None )]
843
+ for node , parent in nodes :
844
+ # Skip the root of the tree, which has parent None
845
+ handler = self .getNodeHandler (node .__class__ ) if parent else False
846
+ if handler and handler not in _may_flatten :
847
+ yield node , parent , handler
848
+ else :
849
+ nodes [:] += ((child , node )
850
+ for child in iter_child_nodes (node ,
851
+ omit ,
852
+ _fields_order ))
853
+
854
+ def handleChildrenFlattened (self , tree , omit = None ):
855
+ """Handle all children recursively as a flat list where possible."""
856
+ for node , parent , handler in self ._iter_flattened (tree , omit = omit ):
857
+ self .handleNode (node , parent , handler )
858
+
826
859
def isLiteralTupleUnpacking (self , node ):
827
860
if isinstance (node , ast .Assign ):
828
861
for child in node .targets + [node .value ]:
@@ -852,7 +885,12 @@ def getDocstring(self, node):
852
885
853
886
return (node .s , doctest_lineno )
854
887
855
- def handleNode (self , node , parent ):
888
+ def handleNode (self , node , parent , handler = None ):
889
+ """
890
+ Handle a single node, invoking its handler, which may recurse.
891
+
892
+ If handler is None, the default handler is used.
893
+ """
856
894
if node is None :
857
895
return
858
896
if self .offset and getattr (node , 'lineno' , None ) is not None :
@@ -863,11 +901,18 @@ def handleNode(self, node, parent):
863
901
if self .futuresAllowed and not (isinstance (node , ast .ImportFrom ) or
864
902
self .isDocstring (node )):
865
903
self .futuresAllowed = False
866
- self . nodeDepth += 1
867
- node .depth = self .nodeDepth
904
+
905
+ node .depth = self .nodeDepth + 1
868
906
node .parent = parent
869
- try :
907
+
908
+ if handler is False :
909
+ return
910
+
911
+ if not handler :
870
912
handler = self .getNodeHandler (node .__class__ )
913
+
914
+ self .nodeDepth += 1
915
+ try :
871
916
handler (node )
872
917
finally :
873
918
self .nodeDepth -= 1
@@ -964,21 +1009,22 @@ def ignore(self, node):
964
1009
pass
965
1010
966
1011
# "stmt" type nodes
967
- DELETE = PRINT = FOR = ASYNCFOR = WHILE = IF = WITH = WITHITEM = \
968
- ASYNCWITH = ASYNCWITHITEM = TRYFINALLY = EXEC = \
969
- EXPR = ASSIGN = handleChildren
1012
+ DELETE = PRINT = EXEC = EXPR = handleChildrenFlattened
1013
+ ASSIGN = TRYFINALLY = handleChildren
1014
+ FOR = ASYNCFOR = WHILE = IF = WITH = ASYNCWITH = handleChildren
1015
+ WITHITEM = ASYNCWITHITEM = handleChildrenFlattened
970
1016
971
1017
PASS = ignore
972
1018
973
1019
# "expr" type nodes
974
1020
BOOLOP = BINOP = UNARYOP = IFEXP = SET = \
975
1021
COMPARE = CALL = REPR = ATTRIBUTE = SUBSCRIPT = \
976
- STARRED = NAMECONSTANT = handleChildren
1022
+ STARRED = NAMECONSTANT = handleChildrenFlattened
977
1023
978
1024
NUM = STR = BYTES = ELLIPSIS = ignore
979
1025
980
1026
# "slice" type nodes
981
- SLICE = EXTSLICE = INDEX = handleChildren
1027
+ SLICE = EXTSLICE = INDEX = handleChildrenFlattened
982
1028
983
1029
# expression contexts are node instances too, though being constants
984
1030
LOAD = STORE = DEL = AUGLOAD = AUGSTORE = PARAM = ignore
@@ -1003,7 +1049,8 @@ def RAISE(self, node):
1003
1049
self .report (messages .RaiseNotImplemented , node )
1004
1050
1005
1051
# additional node types
1006
- COMPREHENSION = KEYWORD = FORMATTEDVALUE = JOINEDSTR = handleChildren
1052
+ COMPREHENSION = handleChildren
1053
+ KEYWORD = FORMATTEDVALUE = JOINEDSTR = handleChildrenFlattened
1007
1054
1008
1055
def DICT (self , node ):
1009
1056
# Complain if there are duplicate keys with different values
@@ -1083,7 +1130,7 @@ def GENERATOREXP(self, node):
1083
1130
self .handleChildren (node )
1084
1131
self .popScope ()
1085
1132
1086
- LISTCOMP = handleChildren if PY2 else GENERATOREXP
1133
+ LISTCOMP = handleChildrenNested if PY2 else GENERATOREXP
1087
1134
1088
1135
DICTCOMP = SETCOMP = GENERATOREXP
1089
1136
0 commit comments