PY-18741 Use type annotations from line comments after function parameters

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)}

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]",