diff --git a/pyt/helper_visitors/vars_visitor.py b/pyt/helper_visitors/vars_visitor.py index 22348528..44366d84 100644 --- a/pyt/helper_visitors/vars_visitor.py +++ b/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) diff --git a/tests/helper_visitors/vars_visitor_test.py b/tests/helper_visitors/vars_visitor_test.py index 849206b8..a354c892 100644 --- a/tests/helper_visitors/vars_visitor_test.py +++ b/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, [])