From b25fc9d6c979831361c50f391078343e71dda992 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Sat, 25 Nov 2017 23:10:37 +0300
Subject: [PATCH] PY-24990 Initial attempt to support generic types for self
 and cls

by including the mapping for these special parameters in PyArgumentsMapping
---
 .../psi/impl/PyCallExpressionHelper.java      | 11 +++-
 .../com/jetbrains/python/PyTypingTest.java    | 64 +++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
index d507b54f48a5..d9d5a3c75a22 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
@@ -802,13 +802,22 @@ public class PyCallExpressionHelper {
       filterExplicitParameters(parameters, callable, callSite, resolveContext);
 
     final ArgumentMappingResults mappingResults = analyzeArguments(arguments, explicitParameters);
+    final Map<PyExpression, PyCallableParameter> mapped = new LinkedHashMap<>();
+    final PyExpression receiver = callSite.getReceiver(callable);
+    if (receiver != null && !explicitParameters.equals(parameters)) {
+      final PyCallableParameter first = ContainerUtil.getFirstItem(parameters);
+      if (first != null && first.getParameter() != null && first.getParameter().isSelf()) {
+        mapped.put(receiver, first);
+      }
+    }
+    mapped.putAll(mappingResults.getMappedParameters());
 
     final PyCallExpression.PyMarkedCallee markedCallee =
       new PyCallExpression.PyMarkedCallee(callableType, callable, null, 0, false, RatedResolveResult.RATE_NORMAL);
 
     return new PyCallExpression.PyArgumentsMapping(callSite,
                                                    markedCallee,
-                                                   mappingResults.getMappedParameters(),
+                                                   mapped,
                                                    mappingResults.getUnmappedParameters(),
                                                    mappingResults.getUnmappedArguments(),
                                                    mappingResults.getParametersMappedToVariadicPositionalArguments(),
diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java
index 2255f76b1c12..0ba62f545737 100644
--- a/python/testSrc/com/jetbrains/python/PyTypingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java
@@ -1083,6 +1083,70 @@ public class PyTypingTest extends PyTestCase {
                              "expr = C().attr");
   }
 
+  // PY-24990
+  public void testSelfAnnotationSameClassInstance() {
+    doTest("C",
+           "from typing import TypeVar\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    def method(self: T) -> T:\n" +
+           "        pass\n" +
+           "\n" +
+           "expr = C().method()");
+  }
+
+  // PY-24990
+  public void testSelfAnnotationSubclassInstance() {
+    doTest("D",
+           "from typing import TypeVar\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    def method(self: T) -> T:\n" +
+           "        pass\n" +
+           "\n" +
+           "class D(C):\n" +
+           "    pass\n" +
+           "\n" +
+           "expr = D().method()");
+  }
+
+  // PY-24990
+  public void testClsAnnotationSameClassInstance() {
+    doTest("C",
+           "from typing import TypeVar, Type\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    @classmethod\n" +
+           "    def factory(cls: Type[T]) -> T:\n" +
+           "        pass\n" +
+           "\n" +
+           "expr = C.factory()");
+  }
+
+  // PY-24990
+  public void testClsAnnotationSubclassInstance() {
+    doTest("D",
+           "from typing import TypeVar, Type\n" +
+           "\n" +
+           "T = TypeVar('T')\n" +
+           "\n" +
+           "class C:\n" +
+           "    @classmethod\n" +
+           "    def factory(cls: 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());