Add unit test cases

pyt/helper_visitors/vars_visitor.py

@@ -25,7 +25,7 @@ def visit_UnaryOp(self, node):
     def visit_Lambda(self, node):
         self.visit(node.body)
 
-    def visit_IfExpr(self, node):
+    def visit_IfExp(self, node):
         self.visit(node.test)
         self.visit(node.body)
         self.visit(node.orelse)

tests/helper_visitors/vars_visitor_test.py

@@ -59,3 +59,79 @@ def test_keyword_numarg(self):
     def test_subscript(self):
         vars = self.perform_vars_on_expression('l[a] = x + y')
         self.assertEqual(vars.result, ['l', 'a', 'x', 'y'])
+
+    def test_visit_boolop(self):
+        # AND operator
+        var1 = self.perform_vars_on_expression('b = x and y')
+        self.assertEqual(var1.result, ['b', 'x', 'y'])
+
+        # OR operator
+        var2 = self.perform_vars_on_expression('b = x or y')
+        self.assertEqual(var2.result, ['b', 'x', 'y'])
+
+    def test_visit_unaryop(self):
+        vars = self.perform_vars_on_expression('a = not b')
+        self.assertEqual(vars.result, ['a', 'b'])
+
+    def test_visit_lambda(self):
+        vars = self.perform_vars_on_expression('f = lambda x: x + 2')
+        self.assertEqual(vars.result, ['f', 'x'])
+
+    def test_visit_set(self):
+        vars = self.perform_vars_on_expression('{a, b, c}')
+        self.assertEqual(vars.result, ['a', 'b', 'c'])
+
+    def test_visit_tuple(self):
+        vars = self.perform_vars_on_expression('(a, b, c)')
+        self.assertEqual(vars.result, ['a', 'b', 'c'])
+
+    def test_visit_list(self):
+        vars = self.perform_vars_on_expression('[a, b, c]')
+        self.assertEqual(vars.result, ['a', 'b', 'c'])
+
+    def test_visit_yield(self):
+        var1 = self.perform_vars_on_expression('yield exp')
+        self.assertEqual(var1.result, ['exp'])
+
+        var2 = self.perform_vars_on_expression('yield from exp')
+        self.assertEqual(var2.result, ['exp'])
+
+    def test_visit_listcomp(self):
+        vars = self.perform_vars_on_expression(
+            '[item for item in coll if cond]')
+        self.assertEqual(vars.result, ['item', 'item', 'coll', 'cond'])
+
+    def test_visit_setcomp(self):
+        vars = self.perform_vars_on_expression('{a for b in d}')
+        self.assertEqual(vars.result, ['a', 'b', 'd'])
+
+    def test_visit_dictcomp(self):
+        vars = self.perform_vars_on_expression('{k1: v1 for (k2, v2) in d}')
+        self.assertEqual(vars.result, ['k1', 'v1', 'k2', 'v2', 'd'])
+
+    def test_visit_compare(self):
+        vars = self.perform_vars_on_expression('a == b')
+        self.assertEqual(vars.result, ['a', 'b'])
+
+    def test_visit_starred(self):
+        vars = self.perform_vars_on_expression('*m = t')
+        self.assertEqual(vars.result, ['m', 't'])
+
+    def test_visit_ifexp(self):
+        vars = self.perform_vars_on_expression('res if test else orelse')
+        self.assertEqual(vars.result, ['test', 'res', 'orelse'])
+
+    def test_visit_subscript(self):
+        # simple slice
+        vars = self.perform_vars_on_expression('foo.bar[lower:upper:step]')
+        self.assertEqual(vars.result, ['foo', 'foo', 'lower', 'upper', 'step'])
+
+        # extended slice
+        vars = self.perform_vars_on_expression('foo[item1:item2, item3]')
+        self.assertEqual(vars.result, ['foo', 'item1', 'item2', 'item3'])
+
+    def test_visit_await(self):
+        vars = self.perform_vars_on_expression("""
+async def bar():
+    await foo()""")
+        self.assertEqual(vars.result, [])