Preparations for some async changes

jedi/evaluate/__init__.py

@@ -88,7 +88,9 @@ from jedi.evaluate.gradual.conversion import try_stub_to_actual_names, \
     try_stubs_to_actual_context_set
 
 
+@debug.increase_indent
 def _execute(context, arguments):
+    debug.dbg('execute: %s %s', context, arguments)
     try:
         func = context.py__call__
     except AttributeError:

jedi/evaluate/base_context.py

@@ -17,6 +17,7 @@ from jedi.common import BaseContextSet, BaseContext
 from jedi.evaluate.helpers import SimpleGetItemNotFound, execute_evaluated
 from jedi.evaluate.utils import safe_property
 from jedi.evaluate.cache import evaluator_as_method_param_cache
+from jedi.cache import memoize_method
 
 
 class HelperContextMixin(object):
@@ -67,6 +68,12 @@ class HelperContextMixin(object):
             return f.filter_name(filters)
         return f.find(filters, attribute_lookup=not search_global)
 
+    def py__await__(self):
+        await_context_set = self.py__getattribute__(u"__await__")
+        if not await_context_set:
+            debug.warning('Tried to run __await__ on context %s', self)
+        return await_context_set.execute_evaluated()
+
     def eval_node(self, node):
         return self.evaluator.eval_element(self, node)
 
@@ -203,11 +210,9 @@ def iterate_contexts(contexts, contextualized_node=None, is_async=False):
     )
 
 
-class ContextWrapper(HelperContextMixin, object):
+class _ContextWrapperBase(HelperContextMixin):
     py__getattribute__ = Context.py__getattribute__
-
+    predefined_names = {}
-    def __init__(self, wrapped_context):
-        self._wrapped_context = wrapped_context
 
     @safe_property
     def name(self):
@@ -225,8 +230,27 @@ class ContextWrapper(HelperContextMixin, object):
         return cls(*args, **kwargs)
 
     def __getattr__(self, name):
+        assert name != '_wrapped_context'
         return getattr(self._wrapped_context, name)
 
+
+class LazyContextWrapper(_ContextWrapperBase):
+    @safe_property
+    @memoize_method
+    def _wrapped_context(self):
+        return self._get_wrapped_context()
+
+    def __repr__(self):
+        return '<%s>' % (self.__class__.__name__)
+
+    def _get_wrapped_context(self):
+        raise NotImplementedError
+
+
+class ContextWrapper(_ContextWrapperBase):
+    def __init__(self, wrapped_context):
+        self._wrapped_context = wrapped_context
+
     def __repr__(self):
         return '%s(%s)' % (self.__class__.__name__, self._wrapped_context)
 

jedi/evaluate/compiled/context.py

@@ -48,6 +48,7 @@ class CompiledObject(Context):
     @CheckAttribute()
     def py__call__(self, arguments):
         if self.tree_node is not None and self.tree_node.type == 'funcdef':
+            # TODO delete, this is dead code
             from jedi.evaluate.context.function import FunctionContext
             return FunctionContext(
                 self.evaluator,

jedi/evaluate/context/function.py

@@ -13,7 +13,7 @@ from jedi.evaluate.arguments import AnonymousArguments
 from jedi.evaluate.filters import ParserTreeFilter, FunctionExecutionFilter
 from jedi.evaluate.names import ContextName, AbstractNameDefinition, ParamName
 from jedi.evaluate.base_context import ContextualizedNode, NO_CONTEXTS, \
-    ContextSet, TreeContext, ContextWrapper
+    ContextSet, TreeContext, ContextWrapper, LazyContextWrapper
 from jedi.evaluate.lazy_context import LazyKnownContexts, LazyKnownContext, \
     LazyTreeContext
 from jedi.evaluate.context import iterable
@@ -353,6 +353,7 @@ class FunctionExecutionContext(TreeContext):
             else:
                 if evaluator.environment.version_info < (3, 5):
                     return NO_CONTEXTS
+                #return ContextSet({CoroutineObject(evaluator, self)})
                 async_classes = evaluator.typing_module.py__getattribute__('Coroutine')
                 return_contexts = self.get_return_values()
                 # Only the first generic is relevant.
@@ -367,6 +368,34 @@ class FunctionExecutionContext(TreeContext):
                 return self.get_return_values()
 
 
+class CoroutineObject(LazyContextWrapper):
+    def __init__(self, evaluator, function_execution):
+        self.evaluator = evaluator
+        self._function_execution = function_execution
+
+    def _get_wrapped_context(self):
+        c, = self.evaluator.typing_module.py__getattribute__('Coroutine') \
+            .execute_annotation()
+        return c
+
+    def py__await__(self):
+        return ContextSet({CoroutineWrapper(self.evaluator, self._function_execution)})
+
+
+class CoroutineWrapper(LazyContextWrapper):
+    def __init__(self, evaluator, function_execution):
+        self.evaluator = evaluator
+        self._function_execution = function_execution
+
+    def _get_wrapped_context(self):
+        c, = self.evaluator.typing_module.py__getattribute__('Generator') \
+            .execute_annotation()
+        return c
+
+    def py__stop_iteration_returns(self):
+        return self._function_execution.get_return_values()
+
+
 class OverloadedFunctionContext(FunctionMixin, ContextWrapper):
     def __init__(self, function, overloaded_functions):
         super(OverloadedFunctionContext, self).__init__(function)

jedi/evaluate/finder.py

@@ -48,7 +48,6 @@ class NameFinder(object):
         self._found_predefined_types = None
         self._analysis_errors = analysis_errors
 
-    @debug.increase_indent
     def find(self, filters, attribute_lookup):
         """
         :params bool attribute_lookup: Tell to logic if we're accessing the

jedi/evaluate/syntax_tree.py

@@ -99,10 +99,7 @@ def eval_node(context, element):
             context_set = eval_trailer(context, context_set, trailer)
 
         if had_await:
-            await_context_set = context_set.py__getattribute__(u"__await__")
+            return context_set.py__await__().py__stop_iteration_returns()
-            if not await_context_set:
-                debug.warning('Tried to run py__await__ on context %s', context)
-            return await_context_set.execute_evaluated().py__stop_iteration_returns()
         return context_set
     elif typ in ('testlist_star_expr', 'testlist',):
         # The implicit tuple in statements.

jedi/plugins/stdlib.py

@@ -101,7 +101,6 @@ _NAMEDTUPLE_FIELD_TEMPLATE = '''\
 class StdlibPlugin(BasePlugin):
     def execute(self, callback):
         def wrapper(context, arguments):
-            debug.dbg('execute: %s %s', context, arguments)
             try:
                 obj_name = context.name.string_name
             except AttributeError: