Normalize async functions properly (#182)

Author: cblegare

sphinx_autodoc_typehints.py

@@ -170,13 +170,19 @@ def remove_prefix(text, prefix):
     for i, l in enumerate(sourcelines):
         if l.lstrip().startswith("def"):
             idx = i
+            whitespace_separator = "def"
             break
+        elif l.lstrip().startswith("async def"):
+            idx = i
+            whitespace_separator = "async def"
+            break
+
     else:
         return "\n".join(sourcelines)
     fn_def = sourcelines[idx]
 
     # Get a string representing the amount of leading whitespace
-    whitespace = fn_def.split("def")[0]
+    whitespace = fn_def.split(whitespace_separator)[0]
 
     # Add this leading whitespace to all lines before and after the `def`
     aligned_prefix = [whitespace + remove_prefix(s, whitespace) for s in sourcelines[:idx]]

tests/test_sphinx_autodoc_typehints.py

@@ -12,7 +12,7 @@
 
 from sphinx_autodoc_typehints import (
     format_annotation, get_annotation_args, get_annotation_class_name, get_annotation_module,
-    process_docstring)
+    normalize_source_lines, process_docstring)
 
 T = TypeVar('T')
 U = TypeVar('U', covariant=True)
@@ -526,3 +526,19 @@ class dummy_module.DataClass(x)
         ''')
         expected_contents = expected_contents.format(**format_args).replace('–', '--')
         assert text_contents == expected_contents
+
+
+def test_normalize_source_lines_async_def():
+    source = textwrap.dedent("""
+    async def async_function():
+        class InnerClass:
+            def __init__(self): pass
+    """)
+
+    expected = textwrap.dedent("""
+    async def async_function():
+        class InnerClass:
+            def __init__(self): pass
+    """)
+
+    assert normalize_source_lines(source) == expected