Move predefine_names to context

jedi/inference/__init__.py

@@ -209,7 +209,7 @@ class InferenceState(object):
             if len(name_dicts) > 1:
                 result = NO_VALUES
                 for name_dict in name_dicts:
-                    with helpers.predefine_names(context, if_stmt, name_dict):
+                    with context.predefine_names(if_stmt, name_dict):
                         result |= infer_node(context, element)
                 return result
             else:

jedi/inference/context.py

@@ -1,4 +1,5 @@
 from abc import abstractmethod
+from contextlib import contextmanager
 
 from parso.tree import search_ancestor
 from parso.python.tree import Name
@@ -88,6 +89,15 @@ class AbstractContext(object):
     def py__doc__(self):
         return ''
 
+    @contextmanager
+    def predefine_names(self, flow_scope, dct):
+        predefined = self.predefined_names
+        predefined[flow_scope] = dct
+        try:
+            yield
+        finally:
+            del predefined[flow_scope]
+
     def __repr__(self):
         return '%s(%s)' % (self.__class__.__name__, self._value)
 

jedi/inference/helpers.py

@@ -184,16 +184,6 @@ def get_module_names(module, all_scopes):
     return names
 
 
-@contextmanager
-def predefine_names(value, flow_scope, dct):
-    predefined = value.predefined_names
-    predefined[flow_scope] = dct
-    try:
-        yield
-    finally:
-        del predefined[flow_scope]
-
-
 def is_string(value):
     if value.inference_state.environment.version_info.major == 2:
         str_classes = (unicode, bytes)

jedi/inference/syntax_tree.py

@@ -13,7 +13,6 @@ from jedi.inference.base_value import ValueSet, NO_VALUES, ContextualizedNode, \
 from jedi.inference.lazy_value import LazyTreeValue
 from jedi.inference import compiled
 from jedi.inference import recursion
-from jedi.inference import helpers
 from jedi.inference import analysis
 from jedi.inference import imports
 from jedi.inference import arguments
@@ -323,7 +322,7 @@ def _infer_expr_stmt(context, stmt, seek_name=None):
 
             for lazy_value in ordered:
                 dct = {for_stmt.children[1].value: lazy_value.infer()}
-                with helpers.predefine_names(context, for_stmt, dct):
+                with context.predefine_names(for_stmt, dct):
                     t = context.infer_node(rhs)
                     left = _infer_comparison(context, left, operator, t)
             value_set = left

jedi/inference/value/function.py

@@ -7,7 +7,6 @@ from jedi.inference import compiled
 from jedi.inference import recursion
 from jedi.inference import docstrings
 from jedi.inference import flow_analysis
-from jedi.inference import helpers
 from jedi.inference.signature import TreeSignature
 from jedi.inference.arguments import AnonymousArguments
 from jedi.inference.filters import ParserTreeFilter, FunctionExecutionFilter
@@ -274,7 +273,7 @@ class FunctionExecutionContext(ValueContext, TreeContextMixin):
                 ordered = list(ordered)
                 for lazy_value in ordered:
                     dct = {str(for_stmt.children[1].value): lazy_value.infer()}
-                    with helpers.predefine_names(self, for_stmt, dct):
+                    with self.predefine_names(for_stmt, dct):
                         for yield_in_same_for_stmt in yields:
                             for result in self._get_yield_lazy_value(yield_in_same_for_stmt):
                                 yield result

jedi/inference/value/iterable.py

@@ -31,8 +31,7 @@ from jedi.inference import recursion
 from jedi.inference.lazy_value import LazyKnownValue, LazyKnownValues, \
     LazyTreeValue
 from jedi.inference.helpers import get_int_or_none, is_string, \
-    predefine_names, infer_call_of_leaf, reraise_getitem_errors, \
-    SimpleGetItemNotFound
+    infer_call_of_leaf, reraise_getitem_errors, SimpleGetItemNotFound
 from jedi.inference.utils import safe_property, to_list
 from jedi.inference.cache import inference_state_method_cache
 from jedi.inference.filters import LazyAttributeOverwrite, publish_method
@@ -173,7 +172,7 @@ class ComprehensionMixin(object):
                 parent_context,
                 comp_for,
             )
-            with predefine_names(context, comp_for, dct):
+            with context.predefine_names(comp_for, dct):
                 try:
                     for result in self._nested(comp_fors[1:], context):
                         yield result