From d6fa1ba18c80dc62f8d3542bb8d6f5721c955a49 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Tue, 28 Nov 2017 17:51:54 +0300
Subject: [PATCH] PY-24990 Duplicate respective tests for function type
 comments

---
 .../typing/PyTypingTypeProvider.java          | 17 ++---
 .../com/jetbrains/python/PyTypingTest.java    | 69 +++++++++++++++++++
 2 files changed, 78 insertions(+), 8 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
index 2027723c92f5..d96087e0ebc7 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
@@ -160,18 +160,18 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
       return null;
     }
     final PyParameterTypeList list = annotation.getParameterTypeList();
-    final List<PyExpression> params = list.getParameterTypes();
-    if (params.size() == 1) {
-      final PyNoneLiteralExpression noneExpr = as(params.get(0), PyNoneLiteralExpression.class);
+    final List<PyExpression> paramTypes = list.getParameterTypes();
+    if (paramTypes.size() == 1) {
+      final PyNoneLiteralExpression noneExpr = as(paramTypes.get(0), PyNoneLiteralExpression.class);
       if (noneExpr != null && noneExpr.isEllipsis()) {
         return Ref.create();
       }
     }
-    final int startOffset = omitFirstParamInTypeComment(func) ? 1 : 0;
+    final int startOffset = omitFirstParamInTypeComment(func, annotation) ? 1 : 0;
     final List<PyParameter> funcParams = Arrays.asList(func.getParameterList().getParameters());
     final int i = funcParams.indexOf(param) - startOffset;
-    if (i >= 0 && i < params.size()) {
-      return getParameterTypeFromFunctionComment(params.get(i), context);
+    if (i >= 0 && i < paramTypes.size()) {
+      return getParameterTypeFromFunctionComment(paramTypes.get(i), context);
     }
     return null;
   }
@@ -253,8 +253,9 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     return new PyCustomType(PROTOCOL, null, false);
   }
 
-  private static boolean omitFirstParamInTypeComment(@NotNull PyFunction func) {
-    return func.getContainingClass() != null && func.getModifier() != PyFunction.Modifier.STATICMETHOD;
+  private static boolean omitFirstParamInTypeComment(@NotNull PyFunction func, @NotNull PyFunctionTypeAnnotation annotation) {
+    return func.getContainingClass() != null && func.getModifier() != PyFunction.Modifier.STATICMETHOD &&
+           annotation.getParameterTypeList().getParameterTypes().size() < func.getParameterList().getParameters().length;
   }
 
   @Nullable
diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java
index 0ba62f545737..fccdba981070 100644
--- a/python/testSrc/com/jetbrains/python/PyTypingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java
@@ -25,6 +25,7 @@ import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.types.TypeEvalContext;
+import org.intellij.lang.annotations.Language;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -1147,6 +1148,74 @@ public class PyTypingTest extends PyTestCase {
            "expr = D.factory()");
   }
 
+  // PY-24990
+  public void testSelfAnnotationInTypeCommentSameClassInstance() {
+    doTest("C",
+           "from typing import TypeVar\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    def method(self):\n" +
+           "        # type: (T) -> T\n" +
+           "        pass\n" +
+           "\n" +
+           "expr = C().method()");
+  }
+
+  // PY-24990
+  public void testSelfAnnotationInTypeCommentSubclassInstance() {
+    doTest("D",
+           "from typing import TypeVar\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    def method(self):\n" +
+           "        # type: (T) -> T\n" +
+           "        pass\n" +
+           "\n" +
+           "class D(C):\n" +
+           "    pass\n" +
+           "\n" +
+           "expr = D().method()");
+  }
+
+  // PY-24990
+  public void testClsAnnotationInTypeCommentSameClassInstance() {
+    doTest("C",
+           "from typing import TypeVar, Type\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    @classmethod\n" +
+           "    def factory(cls) -> T:\n" +
+           "        # type: (Type[T]) -> T\n" +
+           "        pass\n" +
+           "\n" +
+           "expr = C.factory()");
+  }
+
+  // PY-24990
+  public void testClsAnnotationInTypeCommentSubclassInstance() {
+    doTest("D",
+           "from typing import TypeVar, Type\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    @classmethod\n" +
+           "    def factory(cls):\n" +
+           "        # type: (Type[T]) -> T\n" +
+           "        pass\n" +
+           "\n" +
+           "class D(C): \n" +
+           "    pass\n" +
+           "\n" +
+           "expr = D.factory()");
+  }
+
   private void doTestNoInjectedText(@NotNull String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());