add support for dictcomp

Author: Lkruitwagen

simpleeval.py

@@ -661,6 +661,7 @@ def __init__(self, operators=None, functions=None, names=None):
                 ast.Set: self._eval_set,
                 ast.ListComp: self._eval_comprehension,
                 ast.GeneratorExp: self._eval_comprehension,
+                ast.DictComp: self._eval_comprehension,
             }
         )
 
@@ -699,7 +700,11 @@ def _eval_set(self, node):
         return set(self._eval(x) for x in node.elts)
 
     def _eval_comprehension(self, node):
-        to_return = []
+
+        if isinstance(node, ast.DictComp):
+            to_return = {}
+        else:
+            to_return = []
 
         extra_names = {}
 
@@ -738,7 +743,10 @@ def do_generator(gi=0):
                     if len(node.generators) > gi + 1:
                         do_generator(gi + 1)
                     else:
-                        to_return.append(self._eval(node.elt))
+                        if isinstance(to_return, dict):
+                            to_return[self._eval(node.key)] = self._eval(node.value)
+                        elif isinstance(to_return, list):
+                            to_return.append(self._eval(node.elt))
 
         try:
             do_generator()

test_simpleeval.py

@@ -729,6 +729,24 @@ def test_unpack(self):
     def test_nested_unpack(self):
         self.t("[a+b+c for a, (b, c) in ((1,(1,1)),(3,(2,2)))]", [3, 7])
 
+    def test_dictcomp_basic(self):
+        self.t("{a:a + 1 for a in [1,2,3]}", {1:2, 2:3, 3:4})
+
+    def test_dictcomp_with_self_reference(self):
+        self.t("{a:a + a for a in [1,2,3]}", {1:2, 2:4, 3:6})
+
+    def test_dictcomp_with_if(self):
+        self.t("{a:a for a in [1,2,3,4,5] if a <= 3}", {1:1, 2:2, 3:3})
+
+    def test_dictcomp_with_multiple_if(self):
+        self.t("{a:a for a in [1,2,3,4,5] if a <= 3 and a > 1 }", {2:2, 3:3})
+
+    def test_dictcomp_unpack(self):
+        self.t("{a:a+b for a,b in ((1,2),(3,4))}", {1:3, 3:7})
+
+    def test_dictcomp_nested_unpack(self):
+        self.t("{a:a+b+c for a, (b, c) in ((1,(1,1)),(3,(2,2)))}", {1:3, 3:7})
+
     def test_other_places(self):
         self.s.functions = {"sum": sum}
         self.t("sum([a+1 for a in [1,2,3,4,5]])", 20)