PY-53120 "Add method" action does not create async def

GitOrigin-RevId: 8b2286a3f50d7b0d00ebbb0befc9811a80f7b304

python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/AddMethodQuickFix.java

@@ -12,10 +12,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
-import com.jetbrains.python.PyNames;
-import com.jetbrains.python.PyPsiBundle;
-import com.jetbrains.python.PythonTemplateRunner;
-import com.jetbrains.python.PythonUiService;
+import com.jetbrains.python.*;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.ParamHelper;
 import com.jetbrains.python.psi.impl.PyFunctionBuilder;
@@ -73,10 +70,14 @@ public class AddMethodQuickFix implements LocalQuickFix {
       PsiElement pe = problemElement.getParent();
       String decoratorName = null; // set to non-null to add a decorator
       PyExpression[] args = PyExpression.EMPTY_ARRAY;
-      if (pe instanceof PyCallExpression) {
-        PyArgumentList arglist = ((PyCallExpression)pe).getArgumentList();
+      if (pe instanceof PyCallExpression callExpression) {
+        PyArgumentList arglist = callExpression.getArgumentList();
         if (arglist == null) return;
         args = arglist.getArguments();
+        if (callExpression.getParent() instanceof PyPrefixExpression prefixExpression &&
+            prefixExpression.getOperator() == PyTokenTypes.AWAIT_KEYWORD) {
+          builder.makeAsync();
+        }
       }
       boolean madeInstance = false;
       if (callByClass) {

python/testData/quickFixes/PyAddMethodQuickFixTest/addAsyncMethodFromInstance.py

@@ -0,0 +1,7 @@
+class A:
+    def __init__(self):
+        self.x = 1
+
+async def f():
+    a = A()
+    await a.<caret><warning descr="Unresolved attribute reference 'y' for class 'A'">y</warning>()

python/testData/quickFixes/PyAddMethodQuickFixTest/addAsyncMethodFromInstance_after.py

@@ -0,0 +1,11 @@
+class A:
+    def __init__(self):
+        self.x = 1
+
+    async def y(self):
+        pass
+
+
+async def f():
+    a = A()
+    await a.y()

python/testData/quickFixes/PyAddMethodQuickFixTest/addAsyncMethodFromMethod.py

@@ -0,0 +1,11 @@
+class A:
+    def __init__(self):
+        self.x = 1
+
+    async def foo(self, a):
+        await self.<caret><warning descr="Unresolved attribute reference 'y' for class 'A'">y</warning>(1, a)
+
+# Some comment
+
+class B:
+    pass

python/testData/quickFixes/PyAddMethodQuickFixTest/addAsyncMethodFromMethod_after.py

@@ -0,0 +1,15 @@
+class A:
+    def __init__(self):
+        self.x = 1
+
+    async def foo(self, a):
+        await self.y(1, a)
+
+    async def y(self, param, a):
+        pass
+
+
+# Some comment
+
+class B:
+    pass

python/testSrc/com/jetbrains/python/quickFixes/PyAddMethodQuickFixTest.java

@@ -43,7 +43,17 @@ public class PyAddMethodQuickFixTest extends PyQuickFixTestCase {
     doQuickFixTest(PyUnresolvedReferencesInspection.class, PyPsiBundle.message("QFIX.add.method.to.class", "y", "A"));
   }
 
+  // PY-53120
+  public void testAddAsyncMethodFromInstance() {
+    doQuickFixTest(PyUnresolvedReferencesInspection.class, PyPsiBundle.message("QFIX.add.method.to.class", "y", "A"));
+  }
+
   public void testAddMethodFromMethod() {
     doQuickFixTest(PyUnresolvedReferencesInspection.class, PyPsiBundle.message("QFIX.add.method.to.class", "y", "A"));
   }
+
+  // PY-53120
+  public void testAddAsyncMethodFromMethod() {
+    doQuickFixTest(PyUnresolvedReferencesInspection.class, PyPsiBundle.message("QFIX.add.method.to.class", "y", "A"));
+  }
 }