133: Support for LHS functions and no comparison while tests

Author: adrianbn

examples/example_inputs/while_func_comparator.py

@@ -1,6 +1,6 @@
 def foo():
- return 6
+ return True
 
-while x < foo():
+while foo():
  print(x)
  x += 1

examples/example_inputs/while_func_comparator_lhs.py

@@ -0,0 +1,6 @@
+def foo():
+ return 6
+
+while foo() > x:
+ print(x)
+ x += 1

examples/example_inputs/while_func_comparator_rhs.py

@@ -0,0 +1,6 @@
+def foo():
+ return 6
+
+while x < foo():
+ print(x)
+ x += 1

pyt/cfg/stmt_visitor.py

@@ -563,19 +563,29 @@ def visit_For(self, node):
 
  def visit_While(self, node):
  label_visitor = LabelVisitor()
- label_visitor.visit(node.test)
+ test = node.test # the test condition of the while loop
+ label_visitor.visit(test)
 
  while_node = self.append_node(Node(
  'while ' + label_visitor.result + ':',
  node,
  path=self.filenames[-1]
  ))
 
- for comp in node.test.comparators:
- if isinstance(comp, ast.Call) and get_call_names_as_string(comp.func) in self.function_names:
- last_node = self.visit(comp)
+ def process_comparator(comp_n):
+ if isinstance(comp_n, ast.Call) and get_call_names_as_string(comp_n.func) in self.function_names:
+ last_node = self.visit(comp_n)
  last_node.connect(while_node)
 
+ if isinstance(test, ast.Compare):
+ comparators = test.comparators
+ comparators.append(test.left) # quirk. See https://greentreesnakes.readthedocs.io/en/latest/nodes.html#Compare
+
+ for comp in comparators:
+ process_comparator(comp)
+ else: # while foo():
+ process_comparator(test)
+
  return self.loop_node_skeleton(while_node, node)
 
  def add_blackbox_or_builtin_call(self, node, blackbox): # noqa: C901

tests/cfg/cfg_test.py

@@ -684,7 +684,7 @@ def test_while_line_numbers(self):
  self.assertLineNumber(else_body_2, 6)
  self.assertLineNumber(next_stmt, 7)
 
- def test_while_func_iterator(self):
+ def test_while_func_comparator(self):
  self.cfg_create_from_file('examples/example_inputs/while_func_comparator.py')
 
  self.assert_length(self.cfg.nodes, expected_length=9)
@@ -699,6 +699,23 @@ def test_while_func_iterator(self):
  body_1 = 7
  _exit = 8
 
+ self.assertEqual(self.cfg.nodes[test].label, 'while foo():')
+
+ def test_while_func_comparator_rhs(self):
+ self.cfg_create_from_file('examples/example_inputs/while_func_comparator_rhs.py')
+
+ self.assert_length(self.cfg.nodes, expected_length=9)
+
+ entry = 0
+ test = 1
+ entry_foo = 2
+ ret_foo = 3
+ exit_foo = 4
+ call_foo = 5
+ _print = 6
+ body_1 = 7
+ _exit = 8
+
  self.assertEqual(self.cfg.nodes[test].label, 'while x < foo():')
 
  self.assertInCfg([
@@ -707,13 +724,41 @@ def test_while_func_iterator(self):
  (_print, test),
  (_exit, test),
  (body_1, _print),
-
  (test, body_1),
  (test, call_foo),
  (ret_foo, entry_foo),
  (exit_foo, ret_foo),
- (call_foo, exit_foo),
+ (call_foo, exit_foo)
+ ])
+
+ def test_while_func_comparator_lhs(self):
+ self.cfg_create_from_file('examples/example_inputs/while_func_comparator_lhs.py')
 
+ self.assert_length(self.cfg.nodes, expected_length=9)
+
+ entry = 0
+ test = 1
+ entry_foo = 2
+ ret_foo = 3
+ exit_foo = 4
+ call_foo = 5
+ _print = 6
+ body_1 = 7
+ _exit = 8
+
+ self.assertEqual(self.cfg.nodes[test].label, 'while foo() > x:')
+
+ self.assertInCfg([
+ (test, entry),
+ (entry_foo, test),
+ (_print, test),
+ (_exit, test),
+ (body_1, _print),
+ (test, body_1),
+ (test, call_foo),
+ (ret_foo, entry_foo),
+ (exit_foo, ret_foo),
+ (call_foo, exit_foo)
  ])