From d884d8452070ce56217cc9a4f8c2cbae1e2fad63 Mon Sep 17 00:00:00 2001
From: Petr Golubev <petr.golubev@jetbrains.com>
Date: Wed, 29 Nov 2023 19:11:50 +0100
Subject: [PATCH] PY-40923 Create function intention does not recognized async
 invocations

GitOrigin-RevId: d1c1c2891ce66c3d2ae7faafc6683e40d772304e
---
 .../quickfix/UnresolvedRefCreateFunctionQuickFix.java    | 9 +++++++++
 .../inspections/UnresolvedRefCreateAsyncFunction.py      | 3 +++
 .../UnresolvedRefCreateAsyncFunction_after.py            | 7 +++++++
 python/testSrc/com/jetbrains/python/PyQuickFixTest.java  | 6 ++++++
 4 files changed, 25 insertions(+)
 create mode 100644 python/testData/inspections/UnresolvedRefCreateAsyncFunction.py
 create mode 100644 python/testData/inspections/UnresolvedRefCreateAsyncFunction_after.py

diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/UnresolvedRefCreateFunctionQuickFix.java b/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/UnresolvedRefCreateFunctionQuickFix.java
index 2058f893aa36..7ba8c7a5c6ca 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/UnresolvedRefCreateFunctionQuickFix.java
+++ b/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);
diff --git a/python/testData/inspections/UnresolvedRefCreateAsyncFunction.py b/python/testData/inspections/UnresolvedRefCreateAsyncFunction.py
new file mode 100644
index 000000000000..884a3b5b671d
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/python/testData/inspections/UnresolvedRefCreateAsyncFunction_after.py b/python/testData/inspections/UnresolvedRefCreateAsyncFunction_after.py
new file mode 100644
index 000000000000..7e725ae5f202
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
index b3becf7d4873..a8922f77b432 100644
--- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
+++ b/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");