@@ -812,9 +812,42 @@ def on_conditional_branch():
812
812
self .report (messages .UndefinedName , node , name )
813
813
814
814
def handleChildren (self , tree , omit = None ):
815
+ """Handle all children recursively, but may be flattened."""
815
816
for node in iter_child_nodes (tree , omit = omit ):
816
817
self .handleNode (node , tree )
817
818
819
+ def handleChildrenNested (self , node ):
820
+ """Handle all children recursively."""
821
+ self .handleChildren (node )
822
+
823
+ def _iter_flattened (self , tree , omit , _fields_order = _FieldsOrder ()):
824
+ """
825
+ Yield child nodes of *node* and their children, with handler.
826
+
827
+ The value yielded is a tuple of the node, its parent and its handler.
828
+ The handler may be False to indicate that no handler and no recursion
829
+ is required as the node is part of a flattened list.
830
+ """
831
+ _may_flatten = (self .handleChildren ,
832
+ self .handleChildrenFlattened )
833
+
834
+ nodes = [(tree , None )]
835
+ for node , parent in nodes :
836
+ # Skip the root of the tree, which has parent None
837
+ handler = self .getNodeHandler (node .__class__ ) if parent else False
838
+ if handler and handler not in _may_flatten :
839
+ yield node , parent , handler
840
+ else :
841
+ nodes [:] += ((child , node )
842
+ for child in iter_child_nodes (node ,
843
+ omit ,
844
+ _fields_order ))
845
+
846
+ def handleChildrenFlattened (self , tree , omit = None ):
847
+ """Handle all children recursively as a flat list where possible."""
848
+ for node , parent , handler in self ._iter_flattened (tree , omit = omit ):
849
+ self .handleNode (node , parent , handler )
850
+
818
851
def isLiteralTupleUnpacking (self , node ):
819
852
if isinstance (node , ast .Assign ):
820
853
for child in node .targets + [node .value ]:
@@ -844,7 +877,12 @@ def getDocstring(self, node):
844
877
845
878
return (node .s , doctest_lineno )
846
879
847
- def handleNode (self , node , parent ):
880
+ def handleNode (self , node , parent , handler = None ):
881
+ """
882
+ Handle a single node, invoking its handler, which may recurse.
883
+
884
+ If handler is None, the default handler is used.
885
+ """
848
886
if node is None :
849
887
return
850
888
if self .offset and getattr (node , 'lineno' , None ) is not None :
@@ -855,11 +893,18 @@ def handleNode(self, node, parent):
855
893
if self .futuresAllowed and not (isinstance (node , ast .ImportFrom ) or
856
894
self .isDocstring (node )):
857
895
self .futuresAllowed = False
858
- self . nodeDepth += 1
859
- node .depth = self .nodeDepth
896
+
897
+ node .depth = self .nodeDepth + 1
860
898
node .parent = parent
861
- try :
899
+
900
+ if handler is False :
901
+ return
902
+
903
+ if not handler :
862
904
handler = self .getNodeHandler (node .__class__ )
905
+
906
+ self .nodeDepth += 1
907
+ try :
863
908
handler (node )
864
909
finally :
865
910
self .nodeDepth -= 1
@@ -911,21 +956,22 @@ def ignore(self, node):
911
956
pass
912
957
913
958
# "stmt" type nodes
914
- DELETE = PRINT = FOR = ASYNCFOR = WHILE = IF = WITH = WITHITEM = \
915
- ASYNCWITH = ASYNCWITHITEM = RAISE = TRYFINALLY = EXEC = \
916
- EXPR = ASSIGN = handleChildren
959
+ DELETE = PRINT = EXEC = EXPR = RAISE = handleChildrenFlattened
960
+ ASSIGN = TRYFINALLY = handleChildren
961
+ FOR = ASYNCFOR = WHILE = IF = WITH = ASYNCWITH = handleChildren
962
+ WITHITEM = ASYNCWITHITEM = handleChildrenFlattened
917
963
918
964
PASS = ignore
919
965
920
966
# "expr" type nodes
921
967
BOOLOP = BINOP = UNARYOP = IFEXP = SET = \
922
968
COMPARE = CALL = REPR = ATTRIBUTE = SUBSCRIPT = \
923
- STARRED = NAMECONSTANT = handleChildren
969
+ STARRED = NAMECONSTANT = handleChildrenFlattened
924
970
925
971
NUM = STR = BYTES = ELLIPSIS = ignore
926
972
927
973
# "slice" type nodes
928
- SLICE = EXTSLICE = INDEX = handleChildren
974
+ SLICE = EXTSLICE = INDEX = handleChildrenFlattened
929
975
930
976
# expression contexts are node instances too, though being constants
931
977
LOAD = STORE = DEL = AUGLOAD = AUGSTORE = PARAM = ignore
@@ -937,7 +983,8 @@ def ignore(self, node):
937
983
MATMULT = ignore
938
984
939
985
# additional node types
940
- COMPREHENSION = KEYWORD = FORMATTEDVALUE = JOINEDSTR = handleChildren
986
+ COMPREHENSION = handleChildren
987
+ KEYWORD = FORMATTEDVALUE = JOINEDSTR = handleChildrenFlattened
941
988
942
989
def DICT (self , node ):
943
990
# Complain if there are duplicate keys with different values
@@ -1017,7 +1064,7 @@ def GENERATOREXP(self, node):
1017
1064
self .handleChildren (node )
1018
1065
self .popScope ()
1019
1066
1020
- LISTCOMP = handleChildren if PY2 else GENERATOREXP
1067
+ LISTCOMP = handleChildrenNested if PY2 else GENERATOREXP
1021
1068
1022
1069
DICTCOMP = SETCOMP = GENERATOREXP
1023
1070
0 commit comments