include-async-functions

Author: Darshan808

IPython/core/guarded_eval.py

@@ -721,7 +721,8 @@ def eval_node(node: Union[ast.AST, None], context: EvaluationContext):
         for child_node in node.body:
             result = eval_node(child_node, context)
         return result
-    if isinstance(node, ast.FunctionDef):
+    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+        is_async = isinstance(node, ast.AsyncFunctionDef)
         func_locals = context.transient_locals.copy()
         func_context = context.replace(transient_locals=func_locals)
         is_property = False
@@ -775,6 +776,7 @@ def dummy_function(*args, **kwargs):
 
         dummy_function.__name__ = node.name
         dummy_function.__node__ = node
+        dummy_function.__is_async__ = is_async
         context.transient_locals[node.name] = dummy_function
         return None
     if isinstance(node, ast.ClassDef):
@@ -792,6 +794,11 @@ def dummy_function(*args, **kwargs):
         dummy_class = type(node.name, bases, class_locals)
         context.transient_locals[node.name] = dummy_class
         return None
+    if isinstance(node, ast.Await):
+        value = eval_node(node.value, context)
+        if hasattr(value, "__awaited_type__"):
+            return value.__awaited_type__
+        return value
     if isinstance(node, ast.Assign):
         return _handle_assign(node, context)
     if isinstance(node, ast.AnnAssign):
@@ -954,9 +961,17 @@ def dummy_function(*args, **kwargs):
                 return overridden_return_type
             return _create_duck_for_heap_type(func)
         else:
-            if hasattr(func, "__inferred_return__"):
-                return func.__inferred_return__
+            inferred_return = getattr(func, "__inferred_return__", None)
             return_type = _eval_return_type(func, node, context)
+            if getattr(func, "__is_async__", False):
+                awaited_type = (
+                    inferred_return if inferred_return is not None else return_type
+                )
+                coroutine_duck = ImpersonatingDuck()
+                coroutine_duck.__awaited_type__ = awaited_type
+                return coroutine_duck
+            if inferred_return:
+                return inferred_return
             if return_type is not NOT_EVALUATED:
                 return return_type
         raise GuardRejection(