Merge pull request #150 from danthedeckie/dictcomp-support

Dictcomp support (rebased #133)

simpleeval.py

@@ -62,6 +62,7 @@
 - edgarrmondragon (Edgar Ramírez-Mondragón) Address Python 3.12+ deprecation warnings
 - cedk (Cédric Krier) <[email protected]> Allow running tests with Werror
 - decorator-factory <[email protected]> More security fixes
+- lkruitwagen (Lucas Kruitwagen) Adding support for dict comprehensions
 
 -------------------------------------
 Basic Usage:
@@ -661,6 +662,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 +701,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 = {}
 
@@ -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)