Create Comprehension.py__iter__.

2015-11-14 20:34:33 +01:00
parent 239f0d7213
commit dc2e52fd7d
4 changed files with 83 additions and 23 deletions
--- a/jedi/parser/tree.py
+++ b/jedi/parser/tree.py
@@ -1423,6 +1423,16 @@ class CompFor(BaseNode):
    type = 'comp_for'
    __slots__ = ()

+    def get_comp_fors(self):
+        yield self
+        last = self.children[-1]
+        while True:
+            if isinstance(last, CompFor):
+                yield last
+            elif not is_node(last, 'comp_if'):
+                break
+            last = last.children[-1]
+
    def is_scope(self):
        return True