Merge the async branch

jedi/evaluate/context/function.py

@@ -63,11 +63,19 @@ class FunctionContext(use_metaclass(CachedMetaClass, TreeContext)):
         """
         Created to be used by inheritance.
         """
-        yield_exprs = get_yield_exprs(self.evaluator, self.tree_node)
-        if yield_exprs:
-            return ContextSet(iterable.Generator(self.evaluator, function_execution))
+        is_coroutine = self.tree_node.parent.type == 'async_stmt'
+        is_generator = bool(get_yield_exprs(self.evaluator, self.tree_node))
+
+        if is_coroutine:
+            if is_generator:
+                return ContextSet(iterable.AsyncGenerator(self.evaluator, function_execution))
+            else:
+                return ContextSet(iterable.Coroutine(self.evaluator, function_execution))
         else:
-            return function_execution.get_return_values()
+            if is_generator:
+                return ContextSet(iterable.Generator(self.evaluator, function_execution))
+            else:
+                return function_execution.get_return_values()
 
     def get_function_execution(self, arguments=None):
         if arguments is None:
@@ -171,7 +179,8 @@ class FunctionExecutionContext(TreeContext):
             yield LazyTreeContext(self, node)
 
     @recursion.execution_recursion_decorator(default=iter([]))
-    def get_yield_lazy_contexts(self):
+    def get_yield_lazy_contexts(self, is_async=False):
+        # TODO: if is_async, wrap yield statements in Awaitable/async_generator_asend
         for_parents = [(y, tree.search_ancestor(y, 'for_stmt', 'funcdef',
                                                 'while_stmt', 'if_stmt'))
                        for y in get_yield_exprs(self.evaluator, self.tree_node)]

jedi/evaluate/context/iterable.py

@@ -39,6 +39,19 @@ from jedi.evaluate.base_context import ContextSet, NO_CONTEXTS, Context, \
     TreeContext, ContextualizedNode
 from jedi.parser_utils import get_comp_fors
 
+try:
+    from types import CoroutineType
+except ImportError:
+    HAS_COROUTINE = False
+else:
+    HAS_COROUTINE = True
+
+try:
+    from types import AsyncGeneratorType
+except ImportError:
+    HAS_ASYNC_GENERATOR = False
+else:
+    HAS_ASYNC_GENERATOR = True
 
 class AbstractIterable(Context):
     builtin_methods = {}
@@ -55,6 +68,83 @@ class AbstractIterable(Context):
         return compiled.CompiledContextName(self, self.array_type)
 
 
+@has_builtin_methods
+class CoroutineMixin(object):
+    array_type = None
+
+    def get_filters(self, search_global, until_position=None, origin_scope=None):
+        gen_obj = compiled.create(self.evaluator, CoroutineType)
+        yield SpecialMethodFilter(self, self.builtin_methods, gen_obj)
+        for filter in gen_obj.get_filters(search_global):
+            yield filter
+
+    def py__bool__(self):
+        return True
+
+    def py__class__(self):
+        gen_obj = compiled.create(self.evaluator, CoroutineType)
+        return gen_obj.py__class__()
+
+    @property
+    def name(self):
+        return compiled.CompiledContextName(self, 'coroutine')
+
+
+class Coroutine(CoroutineMixin, Context):
+    def __init__(self, evaluator, func_execution_context):
+        if not HAS_COROUTINE:
+            raise ImportError("Need python3.5 to support coroutines.")
+        super(Coroutine, self).__init__(evaluator, parent_context=evaluator.BUILTINS)
+        self._func_execution_context = func_execution_context
+
+    def execute_await(self):
+        return self._func_execution_context.get_return_values()
+
+    def __repr__(self):
+        return "<%s of %s>" % (type(self).__name__, self._func_execution_context)
+
+
+@has_builtin_methods
+class AsyncGeneratorMixin(object):
+    array_type = None
+
+    @register_builtin_method('__anext__')
+    def py__anext__(self):
+        return ContextSet.from_sets(lazy_context.infer() for lazy_context in self.py__aiter__())
+
+    def get_filters(self, search_global, until_position=None, origin_scope=None):
+        gen_obj = compiled.create(self.evaluator, AsyncGeneratorType)
+        yield SpecialMethodFilter(self, self.builtin_methods, gen_obj)
+        for filter in gen_obj.get_filters(search_global):
+            yield filter
+
+    def py__bool__(self):
+        return True
+
+    def py__class__(self):
+        gen_obj = compiled.create(self.evaluator, AsyncGeneratorType)
+        return gen_obj.py__class__()
+
+    @property
+    def name(self):
+        return compiled.CompiledContextName(self, 'asyncgenerator')
+
+
+class AsyncGenerator(AsyncGeneratorMixin, Context):
+    """Handling of `yield` functions."""
+    def __init__(self, evaluator, func_execution_context):
+        if not HAS_ASYNC_GENERATOR:
+            raise ImportError("Need python3.6 to support async generators.")
+        super(AsyncGenerator, self).__init__(evaluator, parent_context=evaluator.BUILTINS)
+        self._func_execution_context = func_execution_context
+
+    def py__aiter__(self):
+        return self._func_execution_context.get_yield_values(is_async=True)
+
+    def __repr__(self):
+        return "<%s of %s>" % (type(self).__name__, self._func_execution_context)
+
+
 @has_builtin_methods
 class GeneratorMixin(object):
     array_type = None
@@ -128,17 +218,18 @@ class Comprehension(AbstractIterable):
             cls = ListComprehension
         return cls(evaluator, context, atom)
 
-    def __init__(self, evaluator, defining_context, atom):
+    def __init__(self, evaluator, defining_context, atom, is_async=False):
         super(Comprehension, self).__init__(evaluator)
         self._defining_context = defining_context
         self._atom = atom
 
     def _get_comprehension(self):
+        "return 'a for a in b'"
         # The atom contains a testlist_comp
         return self._atom.children[1]
 
     def _get_comp_for(self):
-        # The atom contains a testlist_comp
+        "return CompFor('for a in b')"
         return self._get_comprehension().children[1]
 
     def _eval_node(self, index=0):
@@ -156,13 +247,17 @@ class Comprehension(AbstractIterable):
 
     def _nested(self, comp_fors, parent_context=None):
         comp_for = comp_fors[0]
-        input_node = comp_for.children[3]
+
+        is_async = 'async' == comp_for.children[comp_for.children.index('for') - 1]
+
+        input_node = comp_for.children[comp_for.children.index('in') + 1]
         parent_context = parent_context or self._defining_context
         input_types = parent_context.eval_node(input_node)
+        # TODO: simulate await if self.is_async
 
         cn = ContextualizedNode(parent_context, input_node)
-        iterated = input_types.iterate(cn)
-        exprlist = comp_for.children[1]
+        iterated = input_types.iterate(cn, is_async=is_async)
+        exprlist = comp_for.children[comp_for.children.index('for') + 1]
         for i, lazy_context in enumerate(iterated):
             types = lazy_context.infer()
             dct = unpack_tuple_to_dict(parent_context, types, exprlist)
@@ -675,7 +770,7 @@ class _ArrayInstance(object):
             for addition in additions:
                 yield addition
 
-    def iterate(self, contextualized_node=None):
+    def iterate(self, contextualized_node=None, is_async=False):
         return self.py__iter__()