diff --git a/jedi/evaluate/__init__.py b/jedi/evaluate/__init__.py index 5c327d25..4e8d1e7a 100644 --- a/jedi/evaluate/__init__.py +++ b/jedi/evaluate/__init__.py @@ -352,15 +352,14 @@ class Evaluator(object): and not(tree.is_node(c[1], 'testlist_comp') and len(c[1].children) > 1): return self.eval_element(c[1]) + try: comp_for = c[1].children[1] except (IndexError, AttributeError): pass else: - if isinstance(comp_for, tree.CompFor): - if atom.children[0] == '{': - # TODO dict/set comprehensions should be working. - return set() + if comp_for.type == 'comp_for' \ + or comp_for == ':' and c[1].children[3].type == 'comp_for': # dict return set([iterable.Comprehension.from_atom(self, atom)]) return set([iterable.Array(self, atom)]) diff --git a/jedi/evaluate/iterable.py b/jedi/evaluate/iterable.py index 64e1dc2c..035c0805 100644 --- a/jedi/evaluate/iterable.py +++ b/jedi/evaluate/iterable.py @@ -101,14 +101,17 @@ class GeneratorMethod(IterableWrapper): class Comprehension(IterableWrapper): @staticmethod def from_atom(evaluator, atom): - mapping = { - '(': GeneratorComprehension, - '[': ListComprehension - } - return mapping[atom.children[0]](evaluator, atom) - - def get_parent_until(self, *args, **kwargs): - return self._atom.get_parent_until(*args, **kwargs) + bracket = atom.children[0] + if bracket == '{': + if atom.children[1].children[1] == ':': + cls = DictComprehension + else: + cls = SetComprehension + elif bracket == '(': + cls = GeneratorComprehension + elif bracket == '[': + cls = ListComprehension + return cls(evaluator, atom) def __init__(self, evaluator, atom): self._evaluator = evaluator @@ -177,6 +180,14 @@ class ArrayMixin(object): def py__class__(self): return compiled.builtin_from_name(self._evaluator, self.type) + @safe_property + def parent(self): + return self._evaluator.BUILTINS + + @property + def name(self): + return FakeSequence(self._evaluator, [], self.type).name + class ListComprehension(Comprehension, ArrayMixin): type = 'list' @@ -185,9 +196,9 @@ class ListComprehension(Comprehension, ArrayMixin): all_types = list(self.py__iter__()) return all_types[index] - @property - def name(self): - return FakeSequence(self._evaluator, [], 'list').name + +class SetComprehension(Comprehension, ArrayMixin): + type = 'set' class GeneratorComprehension(Comprehension, GeneratorMixin): @@ -237,13 +248,6 @@ class Array(IterableWrapper, ArrayMixin): else: return self._evaluator.eval_element(self._items()[index]) - @safe_property - def parent(self): - return self._evaluator.BUILTINS - - def get_parent_until(self): - return self._evaluator.BUILTINS - def __getattr__(self, name): if name not in ['start_pos', 'get_only_subelement', 'parent', 'get_parent_until', 'items']: diff --git a/test/completion/comprehensions.py b/test/completion/comprehensions.py index 9e65946f..5aeb7107 100644 --- a/test/completion/comprehensions.py +++ b/test/completion/comprehensions.py @@ -51,26 +51,6 @@ left #? int() [a for a in {1:'x'}][0] -##? str() -{a-1:b for a,b in {1:'a', 3:1.0}.items()}[0] - -# with a set literal -#? int() -[a for a in {1, 2, 3}][0] - -##? set() -{a for a in range(10)} - -##? int() -[x for x in {a for a in range(10)}][0] - -##? int() -{a for a in range(10)}.pop() - -##? int() -iter({a for a in range(10)}).next() - - # list comprehensions should also work in combination with functions def listen(arg): for x in arg: @@ -81,6 +61,9 @@ listen(['' for x in [1]]) #? ([str for x in []])[0] +# with a set literal +#? int() +[a for a in {1, 2, 3}][0] # ----------------- # nested list comprehensions @@ -140,17 +123,34 @@ left right # ----------------- -# set comprehensions +# dict comprehensions # ----------------- ##? str() {a - 1: b for a, b in {1: 'a', 3: 1.0}.items()}[0] - -#? -{a - 1 for a in [1]} #? int() {a - 1: 3 for a in [1]}[0] +# ----------------- +# set comprehensions +# ----------------- + +#? set() +{a - 1 for a in [1]} + +#? set() +{a for a in range(10)} + +#? int() +[x for x in {a for a in range(10)}][0] + +#? int() +{a for a in range(10)}.pop() + +#? int() +iter({a for a in range(10)}).next() + + # ----------------- # name resolution in comprehensions. # -----------------