diff --git a/simpleeval.py b/simpleeval.py index f6a3ed6..ae434a3 100644 --- a/simpleeval.py +++ b/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() diff --git a/test_simpleeval.py b/test_simpleeval.py index bc83d50..5a918e8 100644 --- a/test_simpleeval.py +++ b/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)