Repository: python/mypy

Issue:
stubgen doesn't generate async function definitions
mypy version: 0.641

Given the following file (named `test.py`):

```python
class Foo:
    async def bar(self):
        return None

async def baz():
    return None
```

Running `stubgen --search-path . test` produces `out/test.pyi` with the following contents:

```python
# Stubs for test (Python 3.7)
#
# NOTE: This dynamically typed stub was automatically generated by stubgen.

class Foo:
    def bar(self): ...

def baz(): ...
```

I would expect the output to be:

```python
# Stubs for test (Python 3.7)
#
# NOTE: This dynamically typed stub was automatically generated by stubgen.

class Foo:
    async def bar(self): ...

async def baz(): ...
```


Discussion:
That's a good catch. Do you want to try creating a PR to fix it?
Yeah, I'll give it a try

Resolving patch:
diff --git a/mypy/stubgen.py b/mypy/stubgen.py
--- a/mypy/stubgen.py
+++ b/mypy/stubgen.py
@@ -450,7 +450,7 @@ def visit_func_def(self, o: FuncDef) -> None:
             return
         if self.is_recorded_name(o.name()):
             return
-        if not self._indent and self._state not in (EMPTY, FUNC):
+        if not self._indent and self._state not in (EMPTY, FUNC) and not o.is_awaitable_coroutine:
             self.add('\n')
         if not self.is_top_level():
             self_inits = find_self_initializers(o)
@@ -458,7 +458,7 @@ def visit_func_def(self, o: FuncDef) -> None:
                 init_code = self.get_init(init, value)
                 if init_code:
                     self.add(init_code)
-        self.add("%sdef %s(" % (self._indent, o.name()))
+        self.add("%s%sdef %s(" % (self._indent, 'async ' if o.is_coroutine else '', o.name()))
         self.record_name(o.name())
         args = []  # type: List[str]
         for i, arg_ in enumerate(o.arguments):
@@ -513,13 +513,36 @@ def visit_decorator(self, o: Decorator) -> None:
         if self.is_private_name(o.func.name()):
             return
         for decorator in o.decorators:
-            if isinstance(decorator, NameExpr) and decorator.name in ('property',
-                                                                      'staticmethod',
-                                                                      'classmethod'):
-                self.add('%s@%s\n' % (self._indent, decorator.name))
-            elif (isinstance(decorator, MemberExpr) and decorator.name == 'setter' and
-                  isinstance(decorator.expr, NameExpr)):
-                self.add('%s@%s.setter\n' % (self._indent, decorator.expr.name))
+            if isinstance(decorator, NameExpr):
+                if decorator.name in ('property',
+                                      'staticmethod',
+                                      'classmethod'):
+                    self.add('%s@%s\n' % (self._indent, decorator.name))
+                elif self.import_tracker.module_for.get(decorator.name) in ('asyncio',
+                                                                            'asyncio.coroutines',
+                                                                            'types'):
+                    self.add_coroutine_decorator(o.func, decorator.name, decorator.name)
+            elif isinstance(decorator, MemberExpr):
+                if decorator.name == 'setter' and isinstance(decorator.expr, NameExpr):
+                    self.add('%s@%s.setter\n' % (self._indent, decorator.expr.name))
+                elif decorator.name == 'coroutine':
+                    if (isinstance(decorator.expr, MemberExpr) and
+                        decorator.expr.name == 'coroutines' and
+                        isinstance(decorator.expr.expr, NameExpr) and
+                            (decorator.expr.expr.name == 'asyncio' or
+                             self.import_tracker.reverse_alias.get(decorator.expr.expr.name) ==
+                                'asyncio')):
+                        self.add_coroutine_decorator(o.func,
+                                                     '%s.coroutines.coroutine' %
+                                                     (decorator.expr.expr.name,),
+                                                     decorator.expr.expr.name)
+                    elif (isinstance(decorator.expr, NameExpr) and
+                          (decorator.expr.name in ('asyncio', 'types') or
+                           self.import_tracker.reverse_alias.get(decorator.expr.name) in
+                            ('asyncio', 'asyncio.coroutines', 'types'))):
+                        self.add_coroutine_decorator(o.func,
+                                                     decorator.expr.name + '.coroutine',
+                                                     decorator.expr.name)
         super().visit_decorator(o)
 
     def visit_class_def(self, o: ClassDef) -> None:
@@ -750,6 +773,13 @@ def add_import_line(self, line: str) -> None:
         if line not in self._import_lines:
             self._import_lines.append(line)
 
+    def add_coroutine_decorator(self, func: FuncDef, name: str, require_name: str) -> None:
+        func.is_awaitable_coroutine = True
+        if not self._indent and self._state not in (EMPTY, FUNC):
+            self.add('\n')
+        self.add('%s@%s\n' % (self._indent, name))
+        self.import_tracker.require_name(require_name)
+
     def output(self) -> str:
         """Return the text for the stub."""
         imports = ''


Test patch:
diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test
--- a/test-data/unit/stubgen.test
+++ b/test-data/unit/stubgen.test
@@ -725,7 +725,7 @@ async def f(a):
 [out]
 from typing import Any
 
-def f(a: Any) -> None: ...
+async def f(a: Any) -> None: ...
 
 [case testInferOptionalOnlyFunc]
 class A:
@@ -919,3 +919,316 @@ x: Callable[[int, int], int]
 from typing import Callable
 
 x: Callable[[int, int], int]
+
+[case testAwaitDef]
+class F:
+    async def f(self):
+        return 1
+
+async def g():
+    return 2
+[out]
+class F:
+    async def f(self): ...
+
+async def g(): ...
+
+[case testCoroutineImportAsyncio]
+import asyncio
+
+class F:
+    @asyncio.coroutine
+    def f(self):
+        return 1
+
+@asyncio.coroutine
+def g():
+    return 2
+
+@asyncio.coroutine
+def h():
+    return 3
+[out]
+import asyncio
+
+class F:
+    @asyncio.coroutine
+    def f(self): ...
+
+@asyncio.coroutine
+def g(): ...
+@asyncio.coroutine
+def h(): ...
+
+[case testCoroutineImportAsyncioCoroutines]
+import asyncio.coroutines
+
+class F:
+    @asyncio.coroutines.coroutine
+    def f(self):
+        return 1
+
+@asyncio.coroutines.coroutine
+def g():
+    return 2
+[out]
+import asyncio.coroutines
+
+class F:
+    @asyncio.coroutines.coroutine
+    def f(self): ...
+
+@asyncio.coroutines.coroutine
+def g(): ...
+
+[case testCoroutineImportAsyncioCoroutinesSub]
+import asyncio
+
+class F:
+    @asyncio.coroutines.coroutine
+    def f(self):
+        return 1
+
+@asyncio.coroutines.coroutine
+def g():
+    return 2
+[out]
+import asyncio
+
+class F:
+    @asyncio.coroutines.coroutine
+    def f(self): ...
+
+@asyncio.coroutines.coroutine
+def g(): ...
+
+[case testCoroutineImportTypes]
+import types
+
+class F:
+    @types.coroutine
+    def f(self):
+        return 1
+
+@types.coroutine
+def g():
+    return 2
+[out]
+import types
+
+class F:
+    @types.coroutine
+    def f(self): ...
+
+@types.coroutine
+def g(): ...
+
+[case testCoroutineFromAsyncioImportCoroutine]
+from asyncio import coroutine
+
+class F:
+    @coroutine
+    def f(self):
+        return 1
+
+@coroutine
+def g():
+    return 2
+[out]
+from asyncio import coroutine
+
+class F:
+    @coroutine
+    def f(self): ...
+
+@coroutine
+def g(): ...
+
+[case testCoroutineFromAsyncioCoroutinesImportCoroutine]
+from asyncio.coroutines import coroutine
+
+class F:
+    @coroutine
+    def f(self):
+        return 1
+
+@coroutine
+def g():
+    return 2
+[out]
+from asyncio.coroutines import coroutine
+
+class F:
+    @coroutine
+    def f(self): ...
+
+@coroutine
+def g(): ...
+
+[case testCoroutineFromTypesImportCoroutine]
+from types import coroutine
+
+class F:
+    @coroutine
+    def f(self):
+        return 1
+
+@coroutine
+def g():
+    return 2
+[out]
+from types import coroutine
+
+class F:
+    @coroutine
+    def f(self): ...
+
+@coroutine
+def g(): ...
+
+[case testCoroutineFromAsyncioImportCoroutineAsC]
+from asyncio import coroutine as c
+
+class F:
+    @c
+    def f(self):
+        return 1
+
+@c
+def g():
+    return 2
+[out]
+from asyncio import coroutine as c
+
+class F:
+    @c
+    def f(self): ...
+
+@c
+def g(): ...
+
+[case testCoroutineFromAsyncioCoroutinesImportCoroutineAsC]
+from asyncio.coroutines import coroutine as c
+
+class F:
+    @c
+    def f(self):
+        return 1
+
+@c
+def g():
+    return 2
+[out]
+from asyncio.coroutines import coroutine as c
+
+class F:
+    @c
+    def f(self): ...
+
+@c
+def g(): ...
+
+[case testCoroutineFromTypesImportCoroutineAsC]
+from types import coroutine as c
+
+class F:
+    @c
+    def f(self):
+        return 1
+
+@c
+def g():
+    return 2
+[out]
+from types import coroutine as c
+
+class F:
+    @c
+    def f(self): ...
+
+@c
+def g(): ...
+
+[case testCoroutineImportAsyncioAsA]
+import asyncio as a
+
+class F:
+    @a.coroutine
+    def f(self):
+        return 1
+
+@a.coroutine
+def g():
+    return 2
+[out]
+import asyncio as a
+
+class F:
+    @a.coroutine
+    def f(self): ...
+
+@a.coroutine
+def g(): ...
+
+[case testCoroutineImportAsyncioCoroutinesAsC]
+import asyncio.coroutines as c
+
+class F:
+    @c.coroutine
+    def f(self):
+        return 1
+
+@c.coroutine
+def g():
+    return 2
+[out]
+import asyncio.coroutines as c
+
+class F:
+    @c.coroutine
+    def f(self): ...
+
+@c.coroutine
+def g(): ...
+
+[case testCoroutineImportAsyncioCoroutinesSubAsA]
+import asyncio as a
+
+class F:
+    @a.coroutines.coroutine
+    def f(self):
+        return 1
+
+@a.coroutines.coroutine
+def g():
+    return 2
+[out]
+import asyncio as a
+
+class F:
+    @a.coroutines.coroutine
+    def f(self): ...
+
+@a.coroutines.coroutine
+def g(): ...
+
+[case testCoroutineImportTypesAsT]
+import types as t
+
+class F:
+    @t.coroutine
+    def f(self):
+        return 1
+
+@t.coroutine
+def g():
+    return 2
+[out]
+import types as t
+
+class F:
+    @t.coroutine
+    def f(self): ...
+
+@t.coroutine
+def g(): ...