Make sure to treat async funcdefs the same way as normal funcdefs

parso/python/diff.py

@@ -47,6 +47,14 @@ def _flows_finished(pgen_grammar, stack):
     return True
 
 
+def _func_or_class_has_suite(node):
+    if node.type == 'decorated':
+        node = node.children[-1]
+    if node.type in ('async_funcdef', 'async_stmt'):
+        node = node.children[-1]
+    return node.type in ('classdef', 'funcdef') and node.children[-1].type == 'suite'
+
+
 def suite_or_file_input_is_valid(pgen_grammar, stack):
     if not _flows_finished(pgen_grammar, stack):
         return False
@@ -511,7 +519,7 @@ class _NodesStack(object):
             # binary search.
             if _get_last_line(node) > until_line:
                 # We can split up functions and classes later.
-                if node.type in ('classdef', 'funcdef') and node.children[-1].type == 'suite':
+                if _func_or_class_has_suite(node):
                     new_nodes.append(node)
                 break
 
@@ -522,9 +530,11 @@ class _NodesStack(object):
 
         last_node = new_nodes[-1]
         line_offset_index = -1
-        if last_node.type in ('classdef', 'funcdef'):
-            suite = last_node.children[-1]
-            if suite.type == 'suite':
+        if _func_or_class_has_suite(last_node):
+            suite = last_node
+            while suite.type != 'suite':
+                suite = suite.children[-1]
+
             suite_tos = _NodesStackNode(suite)
             # Don't need to pass line_offset here, it's already done by the
             # parent.
@@ -538,7 +548,7 @@ class _NodesStack(object):
                 new_tos = recursive_tos
                 line_offset_index = -2
 
-        elif (new_nodes[-1].type in ('error_leaf', 'error_node') or
+        elif (last_node.type in ('error_leaf', 'error_node') or
               _is_flow_node(new_nodes[-1])):
             # Error leafs/nodes don't have a defined start/end. Error
             # nodes might not end with a newline (e.g. if there's an

test/test_diff_parser.py

@@ -517,7 +517,7 @@ def test_end_newline_with_decorator(differ):
             json.l''')
 
     differ.initialize(code)
-    module = differ.parse(code + '\n', copies=0)
+    module = differ.parse(code + '\n', copies=1)
     decorated, endmarker = module.children
     assert decorated.type == 'decorated'
     decorator, func = decorated.children