diff --git a/python/python-parser/src/com/jetbrains/python/PyNames.kt b/python/python-parser/src/com/jetbrains/python/PyNames.kt index e9071f43bcae..503632d93a7c 100644 --- a/python/python-parser/src/com/jetbrains/python/PyNames.kt +++ b/python/python-parser/src/com/jetbrains/python/PyNames.kt @@ -109,6 +109,7 @@ object PyNames { const val FUNCTION: String = "function" const val TYPES_FUNCTION_TYPE: String = "types.FunctionType" + const val TYPES_COROUTINE_TYPE: String = "types.CoroutineType" const val TYPES_METHOD_TYPE: String = "types.UnboundMethodType" const val FUTURE_MODULE: String = "__future__" diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt index fb65e6ddf851..b802c0aba09f 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt @@ -221,7 +221,7 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext() { if (typeRef != null) { // Do not use toAsyncIfNeeded, as it also converts Generators. Here we do not need it. if (callable.isAsync && callable.isAsyncAllowed && !callable.isGenerator) { - return Ref(wrapInCoroutineType(typeRef.get(), callable)) + return Ref(typeRef.get().wrapInCoroutineType(callable)) } return typeRef } @@ -2803,7 +2803,7 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext() { fun toAsyncIfNeeded(function: PyFunction, returnType: PyType?): PyType? { if (function.isAsync && function.isAsyncAllowed) { if (!function.isGenerator) { - return wrapInCoroutineType(returnType, function) + return returnType.wrapInCoroutineType(function) } val desc = GeneratorTypeDescriptor.fromGenerator(returnType) if (desc != null) { @@ -2828,9 +2828,17 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext() { } } - private fun wrapInCoroutineType(returnType: PyType?, resolveAnchor: PsiElement): PyType? { - val coroutine = PyPsiFacade.getInstance(resolveAnchor.project).createClassByQName(COROUTINE, resolveAnchor) - return if (coroutine != null) PyCollectionTypeImpl(coroutine, false, listOf(null, null, returnType)) else null + private fun PyType?.wrapInCoroutineType(anchor: PsiElement): PyType? { + val facade = PyPsiFacade.getInstance(anchor.project) + val targetClass = + facade.createClassByQName(PyNames.TYPES_COROUTINE_TYPE, anchor) + ?: facade.createClassByQName(COROUTINE, anchor) + ?: return null + return PyCollectionTypeImpl( + targetClass, + false, + listOf(null, null, this) + ) } @JvmStatic @@ -2851,18 +2859,15 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext() { @JvmStatic fun unwrapCoroutineReturnType(coroutineType: PyType?): Ref? { - val genericType = PyUtil.`as`(coroutineType, PyCollectionType::class.java) + if (coroutineType !is PyCollectionType) return null + val qName = coroutineType.classQName - if (genericType != null) { - val qName = genericType.classQName + if (AWAITABLE == qName) { + return Ref(coroutineType.elementTypes.getOrNull(0)) + } - if (AWAITABLE == qName) { - return Ref(ContainerUtil.getOrElse(genericType.elementTypes, 0, null)) - } - - if (COROUTINE == qName) { - return Ref(ContainerUtil.getOrElse(genericType.elementTypes, 2, null)) - } + if (qName in arrayOf(COROUTINE, PyNames.TYPES_COROUTINE_TYPE)) { + return Ref(coroutineType.elementTypes.getOrNull(2)) } return null diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/PySuspiciousBooleanConditionInspection.kt b/python/python-psi-impl/src/com/jetbrains/python/inspections/PySuspiciousBooleanConditionInspection.kt index 3d91278e3b4f..5128bcae9c3d 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/inspections/PySuspiciousBooleanConditionInspection.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/PySuspiciousBooleanConditionInspection.kt @@ -86,8 +86,7 @@ class PySuspiciousBooleanConditionInspection : PyInspection() { val type = myTypeEvalContext.getType(this) ?: return // Check if the type is a coroutine - // TODO: use `CoroutineType` instead - if (type is PyClassType && type.classQName == PyTypingTypeProvider.COROUTINE) { + if (type is PyClassType && (type.classQName == PyTypingTypeProvider.COROUTINE || type.classQName == PyNames.TYPES_COROUTINE_TYPE)) { registerProblem(this, PyPsiBundle.message("INSP.suspicious.boolean.condition.coroutine"), PyAddAwaitQuickFix()) } } diff --git a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt index 9bedada9e77e..6f64a24e783c 100644 --- a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt +++ b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt @@ -20,6 +20,7 @@ import com.intellij.util.ProcessingContext import com.intellij.util.containers.ContainerUtil import com.intellij.util.containers.toArray import com.jetbrains.python.BaseReference +import com.jetbrains.python.PyNames import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider import com.jetbrains.python.psi.PyArgumentList import com.jetbrains.python.psi.PyCallExpression @@ -145,7 +146,7 @@ class PyTextFixtureTypeProvider : PyTypeProviderBase() { if (ArrayUtil.contains(qName, "typing.Awaitable", PyTypingTypeProvider.GENERATOR)) { return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 0, null)) } - if (PyTypingTypeProvider.COROUTINE == qName) { + if (PyTypingTypeProvider.COROUTINE == qName || PyNames.TYPES_COROUTINE_TYPE == qName) { return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 2, null)) } } diff --git a/python/testData/intentions/SpecifyTypeInPy3AnnotationsIntentionTest/addsImportsWhenNeeded_after.py b/python/testData/intentions/SpecifyTypeInPy3AnnotationsIntentionTest/addsImportsWhenNeeded_after.py index 9e08252ab8e1..a26aa07dc893 100644 --- a/python/testData/intentions/SpecifyTypeInPy3AnnotationsIntentionTest/addsImportsWhenNeeded_after.py +++ b/python/testData/intentions/SpecifyTypeInPy3AnnotationsIntentionTest/addsImportsWhenNeeded_after.py @@ -1,8 +1,9 @@ -from typing import Any, Coroutine +from types import CoroutineType +from typing import Any async def bar() -> int: return 42 -def foo(x, y) -> Coroutine[Any, Any, int]: +def foo(x, y) -> CoroutineType[Any, Any, int]: return bar() \ No newline at end of file diff --git a/python/testData/quickdoc/AsyncFunctionQuickDoc.html b/python/testData/quickdoc/AsyncFunctionQuickDoc.html index 1f715802134e..ecbd9a60e6b2 100644 --- a/python/testData/quickdoc/AsyncFunctionQuickDoc.html +++ b/python/testData/quickdoc/AsyncFunctionQuickDoc.html @@ -1 +1 @@ -
 AsyncFunctionQuickDoc
async def func() -> Coroutine[Any, Any, None]
\ No newline at end of file +
 AsyncFunctionQuickDoc
async def func() -> CoroutineType[Any, Any, None]
\ No newline at end of file diff --git a/python/testData/quickdoc/AsyncFunctionTooltip.html b/python/testData/quickdoc/AsyncFunctionTooltip.html index a82b5a744d6f..cf5e2daa6759 100644 --- a/python/testData/quickdoc/AsyncFunctionTooltip.html +++ b/python/testData/quickdoc/AsyncFunctionTooltip.html @@ -1 +1 @@ -async def func() -> Coroutine[Any, Any, None] \ No newline at end of file +async def func() -> CoroutineType[Any, Any, None] \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/Py3TypeTest.java b/python/testSrc/com/jetbrains/python/Py3TypeTest.java index f2228baf6269..7a8ce32ecc1f 100644 --- a/python/testSrc/com/jetbrains/python/Py3TypeTest.java +++ b/python/testSrc/com/jetbrains/python/Py3TypeTest.java @@ -531,7 +531,7 @@ public class Py3TypeTest extends PyTestCase { } public void testAsyncDefReturnType() { - doTest("Coroutine[Any, Any, int]", + doTest("CoroutineType[Any, Any, int]", """ async def foo(x): await x @@ -565,6 +565,18 @@ public class Py3TypeTest extends PyTestCase { """); } + public void testAwaitOnTypingCoroutineAnnotation() { + doTest("int", + """ + from typing import Any, Coroutine + + x: Coroutine[Any, Any, int] + + async def bar(): + expr = await x + """); + } + // PY-16987 public void testNoTypeInGoogleDocstringParamAnnotation() { doTest("int", """ @@ -1553,7 +1565,7 @@ public class Py3TypeTest extends PyTestCase { // PY-24067 public void testAsyncFunctionReturnTypeInDocstring() { - doTest("Coroutine[Any, Any, int]", + doTest("CoroutineType[Any, Any, int]", """ async def f(): ""\" @@ -1565,7 +1577,7 @@ public class Py3TypeTest extends PyTestCase { // PY-27518 public void testAsyncFunctionReturnTypeInNumpyDocstring() { - doTest("Coroutine[Any, Any, int]", + doTest("CoroutineType[Any, Any, int]", """ async def f(): ""\" @@ -1592,7 +1604,7 @@ public class Py3TypeTest extends PyTestCase { // PY-26643 public void testReplaceSelfInCoroutine() { - doTest("Coroutine[Any, Any, B]", + doTest("CoroutineType[Any, Any, B]", """ class A: async def foo(self): diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index 8aabd17d9e67..1ca3c51bf672 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java @@ -1152,7 +1152,7 @@ public class PyTypingTest extends PyTestCase { } public void testCoroutineReturnsGenerator() { - doTest("Coroutine[Any, Any, Generator[int, Any, Any]]", + doTest("CoroutineType[Any, Any, Generator[int, Any, Any]]", """ from typing import Generator @@ -6271,7 +6271,7 @@ public class PyTypingTest extends PyTestCase { // PY-36416 public void testReturnTypeOfNonAnnotatedAsyncOverride() { - doTest("Coroutine[Any, Any, str]", """ + doTest("CoroutineType[Any, Any, str]", """ class Base: async def get(self) -> str: ... diff --git a/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java b/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java index 90e8cffaade0..cd5f7184c52d 100644 --- a/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java +++ b/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java @@ -97,7 +97,7 @@ public class PyiTypeTest extends PyTestCase { } public void testCoroutineType() { - doTest("Coroutine[Any, Any, int]"); + doTest("CoroutineType[Any, Any, int]"); } public void testPyiOnPythonPath() {