Feature - DictComp support

simpleeval.py

@@ -57,6 +57,7 @@
 - bozokopic (Bozo Kopic) Memory leak fix
 - daxamin (Dax Amin) Better error for attempting to eval empty string
 - smurfix (Matthias Urlichs) Allow clearing functions / operators / etc completely
+- lkruitwagen (Lucas Kruitwagen) Adding support for dict comprehensions
 
 -------------------------------------
 Basic Usage:
@@ -655,6 +656,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,
             }
         )
 
@@ -693,7 +695,10 @@ 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 = {}
 
@@ -732,7 +737,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

@@ -732,6 +732,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)