@@ -555,15 +555,26 @@ def visit_For(self, node):
555555 path = self .filenames [- 1 ]
556556 ))
557557
558- if isinstance (node .iter , ast .Call ) and get_call_names_as_string (node .iter .func ) in self .function_names :
559- last_node = self .visit (node .iter )
560- last_node .connect (for_node )
558+ self .process_loop_funcs (node .iter , for_node )
561559
562560 return self .loop_node_skeleton (for_node , node )
563561
562+ def process_loop_funcs (self , comp_n , loop_node ):
563+ """
564+ If the loop test node contains function calls, it connects the loop node to the nodes of
565+ those function calls.
566+
567+ :param comp_n: The test node of a loop that may contain functions.
568+ :param loop_node: The loop node itself to connect to the new function nodes if any
569+ :return: None
570+ """
571+ if isinstance (comp_n , ast .Call ) and get_call_names_as_string (comp_n .func ) in self .function_names :
572+ last_node = self .visit (comp_n )
573+ last_node .connect (loop_node )
574+
564575 def visit_While (self , node ):
565576 label_visitor = LabelVisitor ()
566- test = node .test # the test condition of the while loop
577+ test = node .test # the test condition of the while loop
567578 label_visitor .visit (test )
568579
569580 while_node = self .append_node (Node (
@@ -572,19 +583,14 @@ def visit_While(self, node):
572583 path = self .filenames [- 1 ]
573584 ))
574585
575- def process_comparator (comp_n ):
576- if isinstance (comp_n , ast .Call ) and get_call_names_as_string (comp_n .func ) in self .function_names :
577- last_node = self .visit (comp_n )
578- last_node .connect (while_node )
579-
580586 if isinstance (test , ast .Compare ):
581587 comparators = test .comparators
582- comparators .append (test .left ) # quirk. See https://greentreesnakes.readthedocs.io/en/latest/nodes.html#Compare
588+ comparators .append (test .left ) # quirk. See https://greentreesnakes.readthedocs.io/en/latest/nodes.html#Compare
583589
584590 for comp in comparators :
585- process_comparator (comp )
586- else : # while foo():
587- process_comparator (test )
591+ self . process_loop_funcs (comp , while_node )
592+ else : # while foo():
593+ self . process_loop_funcs (test , while_node )
588594
589595 return self .loop_node_skeleton (while_node , node )
590596
0 commit comments