PY-40923 Create function intention does not recognized async invocations

GitOrigin-RevId: d1c1c2891ce66c3d2ae7faafc6683e40d772304e

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

@@ -13,6 +13,7 @@ import com.intellij.psi.PsiLanguageInjectionHost;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyPsiBundle;
+import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonTemplateRunner;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.ParamHelper;
@@ -29,9 +30,13 @@ import static com.jetbrains.python.psi.PyUtil.as;
  */
 public class UnresolvedRefCreateFunctionQuickFix implements LocalQuickFix {
   private final String myFunctionName;
+  private final boolean myAsync;
 
   public UnresolvedRefCreateFunctionQuickFix(@NotNull PyReferenceExpression reference) {
     myFunctionName = reference.getReferencedName();
+    PyCallSiteExpression callSiteExpression = as(reference.getParent(), PyCallSiteExpression.class);
+    PyPrefixExpression prefixExpression = callSiteExpression != null ? as(callSiteExpression.getParent(), PyPrefixExpression.class) : null;
+    myAsync = prefixExpression != null && prefixExpression.getOperator() == PyTokenTypes.AWAIT_KEYWORD;
   }
 
   @Nls
@@ -77,6 +82,10 @@ public class UnresolvedRefCreateFunctionQuickFix implements LocalQuickFix {
       functionBuilder.parameter("args");
     }
 
+    if (myAsync) {
+      functionBuilder.makeAsync();
+    }
+
     PyFunction function = functionBuilder.buildFunction();
 
     final InjectedLanguageManager instance = InjectedLanguageManager.getInstance(project);

python/testData/inspections/UnresolvedRefCreateAsyncFunction.py

@@ -0,0 +1,3 @@
+async def f():
+    s = 3
+    await <error descr="Unresolved reference 'ref'"><caret>ref</error>(s, t=1)

python/testData/inspections/UnresolvedRefCreateAsyncFunction_after.py

@@ -0,0 +1,7 @@
+async def ref(s, t):
+    pass
+
+
+async def f():
+    s = 3
+    await ref(s, t=1)

python/testSrc/com/jetbrains/python/PyQuickFixTest.java

@@ -401,6 +401,12 @@ public class PyQuickFixTest extends PyTestCase {
                      PyPsiBundle.message("QFIX.NAME.unresolved.reference.create.function", "ref"), true, true);
   }
 
+  public void testUnresolvedRefCreateAsyncFunction() {
+    runWithLanguageLevel(LanguageLevel.getLatest(),
+                         () -> doInspectionTest(PyUnresolvedReferencesInspection.class,
+                                                PyPsiBundle.message("QFIX.NAME.unresolved.reference.create.function", "ref"), true, true));
+  }
+
   public void testUnresolvedRefNoCreateFunction() {
     myFixture.enableInspections(PyUnresolvedReferencesInspection.class);
     myFixture.configureByFile("UnresolvedRefNoCreateFunction.py");