diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java index e90006f469d6..540c873ea284 100644 --- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java @@ -24,6 +24,7 @@ import com.intellij.openapi.util.Ref; import com.intellij.psi.PsiComment; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiPolyVariantReference; +import com.intellij.psi.PsiWhiteSpace; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.containers.HashSet; import com.jetbrains.python.PyNames; @@ -90,6 +91,14 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { } } } + + if (context.maySwitchToAST(param)) { + final String comment = getTypeComment(param); + if (comment != null) { + return Ref.create(getStringBasedType(comment, param, new Context(context))); + } + } + final String comment = func.getTypeCommentAnnotation(); if (comment != null) { final PyTypeParser.ParseResult result = PyTypeParser.parsePep484FunctionTypeComment(param, comment); @@ -197,6 +206,21 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { return null; } + @Nullable + private static String getTypeComment(@NotNull PyParameter parameter) { + for (PsiElement next = parameter.getNextSibling(); next != null; next = next.getNextSibling()) { + if (next.textContains('\n')) break; + if (!(next instanceof PsiWhiteSpace)) { + if (",".equals(next.getText())) continue; + if (next instanceof PsiComment) { + return getTypeCommentValue(next.getText()); + } + break; + } + } + return null; + } + /** * Checks that text of a comment starts with the "type:" prefix and returns trimmed part afterwards. This trailing part is supposed to * contain type annotation in PEP 484 compatible format, that can be parsed with either {@link PyTypeParser#parse(PsiElement, String)} diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index 96346cdb1433..3f0de7bfce69 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java @@ -514,6 +514,18 @@ public class PyTypingTest extends PyTestCase { "expr = f"); } + // PY-18741 + public void testFunctionTypeCommentWithParamTypeComment() { + doTest("(x: int, y: bool, z: Any) -> str", + "def f(x, # type: int \n" + + " y # type: bool\n" + + " ,z):\n" + + " # type: (...) -> str\n" + + " pass\n" + + "\n" + + "expr = f"); + } + // PY-18386 public void testRecursiveType() { doTest("Union[int, Any]",