Move iterate logic to the context.

jedi/evaluate/context.py

@@ -1,7 +1,7 @@
 from parso.python.tree import ExprStmt, CompFor
 
 from jedi import debug
-from jedi._compatibility import Python3Method
+from jedi._compatibility import Python3Method, zip_longest
 from jedi.parser_utils import clean_scope_docstring, get_doc_with_call_signature
 from jedi.common import BaseContextSet
 
@@ -66,6 +66,22 @@ class Context(object):
 
         return self.evaluator.execute(self, arguments)
 
+    def iterate(self, contextualized_node=None):
+        debug.dbg('iterate')
+        try:
+            iter_method = self.py__iter__
+        except AttributeError:
+            if contextualized_node is not None:
+                from jedi.evaluate import analysis
+                analysis.add(
+                    contextualized_node.context,
+                    'type-error-not-iterable',
+                    contextualized_node._node,
+                    message="TypeError: '%s' object is not iterable" % self)
+            return iter([])
+        else:
+            return iter_method()
+
     def execute_evaluated(self, *value_list):
         """
         Execute a function with already executed arguments.
@@ -248,6 +264,13 @@ class ContextSet(BaseContextSet):
     def py__class__(self):
         return ContextSet.from_iterable(c.py__class__() for c in self._set)
 
+    def iterate(self, contextualized_node=None):
+        type_iters = [c.iterate(contextualized_node) for c in self._set]
+        for lazy_contexts in zip_longest(*type_iters):
+            yield get_merged_lazy_context(
+                [l for l in lazy_contexts if l is not None]
+            )
+
 
 NO_CONTEXTS = ContextSet()
 

jedi/evaluate/iterable.py

@@ -23,7 +23,7 @@ It is important to note that:
 from jedi import debug
 from jedi import settings
 from jedi.evaluate.utils import safe_property
-from jedi._compatibility import unicode, zip_longest, is_py3
+from jedi._compatibility import unicode, is_py3
 from jedi.evaluate.utils import to_list
 from jedi.evaluate import compiled
 from jedi.evaluate import helpers
@@ -240,14 +240,13 @@ class Comprehension(AbstractSequence):
         return CompForContext.from_comp_for(parent_context, comp_for)
 
     def _nested(self, comp_fors, parent_context=None):
-        evaluator = self.evaluator
         comp_for = comp_fors[0]
         input_node = comp_for.children[3]
         parent_context = parent_context or self._defining_context
         input_types = parent_context.eval_node(input_node)
 
         cn = context.ContextualizedNode(parent_context, input_node)
-        iterated = py__iter__(evaluator, input_types, cn)
+        iterated = input_types.iterate(cn)
         exprlist = comp_for.children[1]
         for i, lazy_context in enumerate(iterated):
             types = lazy_context.infer()
@@ -587,7 +586,7 @@ def unpack_tuple_to_dict(context, types, exprlist):
         dct = {}
         parts = iter(exprlist.children[::2])
         n = 0
-        for lazy_context in py__iter__(context.evaluator, types, exprlist):
+        for lazy_context in types.iterate(exprlist):
             n += 1
             try:
                 part = next(parts)
@@ -614,36 +613,14 @@ def unpack_tuple_to_dict(context, types, exprlist):
     raise NotImplementedError
 
 
-def py__iter__(evaluator, types, contextualized_node=None):
+def py__iter__types(evaluator, contexts, contextualized_node=None):
-    debug.dbg('py__iter__')
-    type_iters = []
-    for typ in types:
-        try:
-            iter_method = typ.py__iter__
-        except AttributeError:
-            if contextualized_node is not None:
-                analysis.add(
-                    contextualized_node.context,
-                    'type-error-not-iterable',
-                    contextualized_node._node,
-                    message="TypeError: '%s' object is not iterable" % typ)
-        else:
-            type_iters.append(iter_method())
-
-    for lazy_contexts in zip_longest(*type_iters):
-        yield context.get_merged_lazy_context(
-            [l for l in lazy_contexts if l is not None]
-        )
-
-
-def py__iter__types(evaluator, types, contextualized_node=None):
     """
     Calls `py__iter__`, but ignores the ordering in the end and just returns
     all types that it contains.
     """
     return ContextSet.from_sets(
         lazy_context.infer()
-        for lazy_context in py__iter__(evaluator, types, contextualized_node)
+        for lazy_context in contexts.iterate(contextualized_node)
     )
 
 
@@ -682,7 +659,8 @@ def py__getitem__(evaluator, context, types, trailer):
                 if isinstance(typ, AbstractSequence) and typ.array_type == 'dict':
                     types.remove(typ)
                     result |= typ.dict_values()
-            return result | py__iter__types(evaluator, types)
+            cs = ContextSet.from_set(types)
+            return result | py__iter__types(evaluator, cs)
 
         for typ in types:
             # The actual getitem call.
@@ -739,7 +717,7 @@ def _check_array_additions(context, sequence):
                 result.add(whatever)
         elif add_name in ['extend', 'update']:
             for key, lazy_context in params:
-                result |= set(py__iter__(context.evaluator, lazy_context.infer()))
+                result |= set(lazy_context.infer().iterate())
         return result
 
     temp_param_add, settings.dynamic_params_for_other_modules = \
@@ -827,7 +805,7 @@ class _ArrayInstance(object):
         except StopIteration:
             pass
         else:
-            for lazy in py__iter__(self.instance.evaluator, lazy_context.infer()):
+            for lazy in lazy_context.infer().iterate():
                 yield lazy
 
         from jedi.evaluate import param
@@ -836,6 +814,9 @@ class _ArrayInstance(object):
             for addition in additions:
                 yield addition
 
+    def iterate(self, contextualized_node=None):
+        return self.py__iter__()
+
 
 class Slice(context.Context):
     def __init__(self, context, start, stop, step):

jedi/evaluate/representation.py

@@ -360,7 +360,7 @@ class FunctionExecutionContext(context.TreeContext):
         node = yield_expr.children[1]
         if node.type == 'yield_arg':  # It must be a yield from.
             cn = ContextualizedNode(self, node.children[1])
-            for lazy_context in iterable.py__iter__(self.evaluator, cn.infer(), cn):
+            for lazy_context in cn.infer().iterate(cn):
                 yield lazy_context
         else:
             yield context.LazyTreeContext(self, node)
@@ -395,7 +395,6 @@ class FunctionExecutionContext(context.TreeContext):
                 return
             last_for_stmt = for_stmt
 
-        evaluator = self.evaluator
         for for_stmt, yields in yields_order:
             if for_stmt is None:
                 # No for_stmt, just normal yields.
@@ -405,7 +404,7 @@ class FunctionExecutionContext(context.TreeContext):
             else:
                 input_node = for_stmt.get_testlist()
                 cn = ContextualizedNode(self, input_node)
-                ordered = iterable.py__iter__(evaluator, cn.infer(), cn)
+                ordered = cn.infer().iterate(cn)
                 ordered = list(ordered)
                 for lazy_context in ordered:
                     dct = {str(for_stmt.children[1].value): lazy_context.infer()}

jedi/evaluate/stdlib.py

@@ -12,7 +12,6 @@ compiled module that returns the types for C-builtins.
 import collections
 import re
 
-from jedi.evaluate.utils import unite
 from jedi.evaluate import compiled
 from jedi.evaluate import representation as er
 from jedi.evaluate.instance import InstanceFunctionExecution, \
@@ -200,7 +199,7 @@ def builtins_reversed(evaluator, sequences, obj, arguments):
     if isinstance(lazy_context, LazyTreeContext):
         # TODO access private
         cn = ContextualizedNode(lazy_context._context, lazy_context.data)
-    ordered = list(iterable.py__iter__(evaluator, sequences, cn))
+    ordered = list(sequences.iterate(cn))
 
     rev = list(reversed(ordered))
     # Repack iterator values and then run it the normal way. This is
@@ -232,9 +231,9 @@ def builtins_isinstance(evaluator, objects, types, arguments):
             elif cls_or_tup.name.string_name == 'tuple' \
                     and cls_or_tup.get_root_context() == evaluator.BUILTINS:
                 # Check for tuples.
-                classes = unite(
+                classes = ContextSet.from_sets(
                     lazy_context.infer()
-                    for lazy_context in cls_or_tup.py__iter__()
+                    for lazy_context in cls_or_tup.iterate()
                 )
                 bool_results.add(any(cls in mro for cls in classes))
             else:

jedi/evaluate/syntax_tree.py

@@ -243,8 +243,7 @@ def _eval_expr_stmt(context, stmt, seek_name=None):
             # predictable. Also only do it, if the variable is not a tuple.
             node = for_stmt.get_testlist()
             cn = ContextualizedNode(context, node)
-            from jedi.evaluate import iterable
+            ordered = list(cn.infer().iterate(cn))
-            ordered = list(iterable.py__iter__(context.evaluator, cn.infer(), cn))
 
             for lazy_context in ordered:
                 dct = {for_stmt.children[1].value: lazy_context.infer()}
@@ -551,8 +550,7 @@ def check_tuple_assignments(evaluator, contextualized_name, context_set):
     lazy_context = None
     for index, node in contextualized_name.assignment_indexes():
         cn = ContextualizedNode(contextualized_name.context, node)
-        from jedi.evaluate import iterable
+        iterated = context_set.iterate(cn)
-        iterated = iterable.py__iter__(evaluator, context_set, cn)
         for _ in range(index + 1):
             try:
                 lazy_context = next(iterated)

jedi/evaluate/sys_path.py

@@ -120,10 +120,9 @@ def _paths_from_assignment(module_context, expr_stmt):
         except AssertionError:
             continue
 
-        from jedi.evaluate.iterable import py__iter__
         from jedi.evaluate.syntax_tree import is_string
         cn = ContextualizedNode(module_context.create_context(expr_stmt), expr_stmt)
-        for lazy_context in py__iter__(module_context.evaluator, cn.infer(), cn):
+        for lazy_context in cn.infer().iterate(cn):
             for context in lazy_context.infer():
                 if is_string(context):
                     yield context.obj