Write a CompForContext that is still not in good shape but working

jedi/inference/__init__.py

@@ -79,7 +79,7 @@ from jedi.inference.base_value import ContextualizedName, ContextualizedNode, \
     ValueSet, NO_VALUES, iterate_values
 from jedi.inference.value import ClassValue, FunctionValue, \
     AnonymousInstance, BoundMethod
-from jedi.inference.value.iterable import CompForValue
+from jedi.inference.context import CompForContext
 from jedi.inference.syntax_tree import infer_trailer, infer_expr_stmt, \
     infer_node, check_tuple_assignments
 from jedi.plugins import plugin_manager
@@ -151,7 +151,7 @@ class InferenceState(object):
         return self.project._get_sys_path(self, environment=self.environment, **kwargs)
 
     def infer_element(self, context, element):
-        if isinstance(context, CompForValue):
+        if isinstance(context, CompForContext):
             return infer_node(context, element)
 
         if_stmt = element
@@ -407,7 +407,7 @@ class InferenceState(object):
             elif scope_node.type in ('comp_for', 'sync_comp_for'):
                 if node.start_pos >= scope_node.children[-1].start_pos:
                     return parent_context
-                return CompForValue.from_comp_for(parent_context, scope_node)
+                return CompForContext(parent_context, scope_node)
             raise Exception("There's a scope that was not managed.")
 
         base_node = base_context.tree_node

jedi/inference/context.py

@@ -137,3 +137,30 @@ class ClassContext(AbstractContext):
             until_position=until_position,
             origin_scope=origin_scope
         )
+
+
+class CompForContext(AbstractContext):
+    def __init__(self, parent_context, comp_for):
+        self._parent_context = parent_context
+        self.inference_state = parent_context.inference_state
+        self._tree_node = comp_for
+
+    @property
+    def parent_context(self):
+        return self._parent_context
+
+    def get_root_context(self):
+        return self._parent_context.get_root_context()
+
+    def is_instance(self):
+        return False
+
+    def is_compiled(self):
+        return False
+
+    @property
+    def tree_node(self):
+        return self._tree_node
+
+    def get_filters(self, until_position=None, origin_scope=None):
+        yield ParserTreeFilter(self)

jedi/inference/value/iterable.py

@@ -35,11 +35,11 @@ from jedi.inference.helpers import get_int_or_none, is_string, \
     SimpleGetItemNotFound
 from jedi.inference.utils import safe_property, to_list
 from jedi.inference.cache import inference_state_method_cache
-from jedi.inference.filters import ParserTreeFilter, LazyAttributeOverwrite, \
-    publish_method
+from jedi.inference.filters import LazyAttributeOverwrite, publish_method
 from jedi.inference.base_value import ValueSet, Value, NO_VALUES, \
-    TreeValue, ContextualizedNode, iterate_values, HelperValueMixin, _sentinel
+    ContextualizedNode, iterate_values, HelperValueMixin, _sentinel
 from jedi.parser_utils import get_sync_comp_fors
+from jedi.inference.context import CompForContext
 
 
 class IterableMixin(object):
@@ -111,15 +111,6 @@ class Generator(GeneratorBase):
         return "<%s of %s>" % (type(self).__name__, self._func_execution_context)
 
 
-class CompForValue(TreeValue):
-    @classmethod
-    def from_comp_for(cls, parent_context, comp_for):
-        return cls(parent_context.inference_state, parent_context, comp_for)
-
-    def get_filters(self, until_position=None, origin_scope=None):
-        yield ParserTreeFilter(self)
-
-
 def comprehension_from_atom(inference_state, value, atom):
     bracket = atom.children[0]
     test_list_comp = atom.children[1]
@@ -159,7 +150,7 @@ def comprehension_from_atom(inference_state, value, atom):
 class ComprehensionMixin(object):
     @inference_state_method_cache()
     def _get_comp_for_value(self, parent_context, comp_for):
-        return CompForValue.from_comp_for(parent_context, comp_for)
+        return CompForContext(parent_context, comp_for)
 
     def _nested(self, comp_fors, parent_context=None):
         comp_for = comp_fors[0]