Fix get_context indentation for async functions (#2107)

AUTHORS.txt

@@ -67,6 +67,7 @@ Code Contributors
 - TingJia Wu (@WutingjiaX) <wutingjia@bytedance.com>
 - Nguyễn Hồng Quân <ng.hong.quan@gmail.com>
 - haoran3160-afk (@haoran3160-afk)
+- Eric3-jp (@Eric3-jp) (with OpenAI Codex assistance)
 
 And a few more "anonymous" contributors.
 

jedi/api/__init__.py

@@ -502,6 +502,8 @@ class Script:
             tree_name = name.tree_name
             if tree_name is not None:  # Happens with lambdas.
                 scope = tree_name.get_definition()
+                if scope.parent.type in ('async_stmt', 'async_funcdef'):
+                    scope = scope.parent
                 if scope.start_pos[1] < column:
                     break
             definition = definition.parent()

test/test_api/test_context.py

@@ -53,6 +53,21 @@ def x():
     ]
 '''
 
+async_code = '''\
+async def coro():
+    return None
+'''
+async_method = '''\
+class C:
+    async def coro(self):
+        return None
+'''
+async_nested = '''\
+async def outer():
+    async def inner():
+        return None
+'''
+
 
 @pytest.mark.parametrize(
     'code, line, column, full_name, expected_parents', [
@@ -106,6 +121,21 @@ def x():
         (with_brackets, 3, None, 'myfile', []),
         (with_brackets, 4, 4, 'myfile.x', ['x']),
         (with_brackets, 4, 5, 'myfile.x', ['x']),
+
+        (async_code, 2, 0, 'myfile', []),
+        (async_code, 2, 4, 'myfile.coro', ['coro']),
+        (async_code, 2, 5, 'myfile.coro', ['coro']),
+        (async_code, 2, None, 'myfile.coro', ['coro']),
+        ('@decorator\n' + async_code, 3, 0, 'myfile', []),
+        ('@decorator\n' + async_code, 3, 4, 'myfile.coro', ['coro']),
+        ('@decorator\n' + async_code, 3, 5, 'myfile.coro', ['coro']),
+        (async_method, 3, 4, 'myfile.C', ['C']),
+        (async_method, 3, 8, 'myfile.C.coro', ['C', 'coro']),
+        (async_method, 3, 9, 'myfile.C.coro', ['C', 'coro']),
+        (async_nested, 3, 0, 'myfile', []),
+        (async_nested, 3, 4, 'myfile.outer', ['outer']),
+        (async_nested, 3, 8, None, ['outer', 'inner']),
+        (async_nested, 3, 9, None, ['outer', 'inner']),
     ]
 )
 def test_context(Script, code, line, column, full_name, expected_parents):