Merge pull request #144 from bcaller/bc-py35-dict

Vars visitor handles python 3.5 dict syntax
python-security · Jul 13, 2018 · c7a932f · c7a932f
2 parents a2c4615 + 6b3f155
commit c7a932f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/pyt/helper_visitors/vars_visitor.py b/pyt/helper_visitors/vars_visitor.py
@@ -32,7 +32,8 @@ def visit_IfExp(self, node):
 
     def visit_Dict(self, node):
         for k in node.keys:
-            self.visit(k)
+            if k is not None:
+                self.visit(k)
         for v in node.values:
             self.visit(v)
 

diff --git a/tests/helper_visitors/vars_visitor_test.py b/tests/helper_visitors/vars_visitor_test.py
@@ -136,3 +136,7 @@ async def bar():
                 await foo()
         """.lstrip())
         self.assertEqual(vars.result, [])
+
+    def test_visit_dict(self):
+        vars = self.perform_vars_on_expression('a = {k1: v1, k2: v2, **d1, **d2}')
+        self.assertEqual(vars.result, ['a', 'k1', 'k2', 'v1', 'v2', 'd1', 'd2'])