From 695bb684fff24af7eab47505f8438261911f109b Mon Sep 17 00:00:00 2001
From: Sergey Kovalev <s_kovalev@wargaming.net>
Date: Tue, 11 Jun 2019 13:47:13 +0300
Subject: [PATCH] Add support for dict comprehension

---
 simpleeval.py      | 10 +++++++---
 test_simpleeval.py | 10 ++++++++++
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/simpleeval.py b/simpleeval.py
index 8724a99..857061e 100644
--- a/simpleeval.py
+++ b/simpleeval.py
@@ -524,6 +524,7 @@ def __init__(self, operators=None, functions=None, names=None):
             ast.List: self._eval_list,
             ast.Set: self._eval_set,
             ast.ListComp: self._eval_comprehension,
+            ast.DictComp: lambda node: self._eval_comprehension(node, is_dict=True),
             ast.GeneratorExp: self._eval_comprehension,
         })
 
@@ -544,8 +545,8 @@ def _eval_list(self, node):
     def _eval_set(self, node):
         return set(self._eval(x) for x in node.elts)
 
-    def _eval_comprehension(self, node):
-        to_return = []
+    def _eval_comprehension(self, node, is_dict=False):
+        to_return = {} if is_dict else []
 
         extra_names = {}
 
@@ -584,7 +585,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 is_dict:
+                            to_return[self._eval(node.key)] = self._eval(node.value)
+                        else:
+                            to_return.append(self._eval(node.elt))
 
         try:
             do_generator()
diff --git a/test_simpleeval.py b/test_simpleeval.py
index cdf8cb8..51542f5 100644
--- a/test_simpleeval.py
+++ b/test_simpleeval.py
@@ -650,6 +650,16 @@ def test_no_leaking_names(self):
             self.s.eval('x')
 
 
+class TestDictComprehensions(DRYTest):
+    """ Test the comprehensions support of the compound-types edition of the class. """
+
+    def setUp(self):
+        self.s = EvalWithCompoundTypes()
+
+    def test_basic(self):
+        self.t('{a: a for a in [1,2,3]}', {1: 1, 2: 2, 3: 3})
+
+
 class TestNames(DRYTest):
     """ 'names', what other languages call variables... """