Get Set comprehensions working.

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)])
 

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']:

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.
 # -----------------