PY-22971 Fixed: Support @typing.overload in regular Python files, not only in Python stubs

Add tests for PyChangeSignatureHandler.
Disable change signature on overloads.

python/src/com/jetbrains/python/refactoring/changeSignature/PyChangeSignatureHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2014 JetBrains s.r.o.
+ * Copyright 2000-2017 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -35,6 +35,8 @@ import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.search.PySuperMethodsSearch;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -55,9 +57,17 @@ public class PyChangeSignatureHandler implements ChangeSignatureHandler {
   public PsiElement findTargetMember(@Nullable PsiElement element) {
     final PyCallExpression callExpression = PsiTreeUtil.getParentOfType(element, PyCallExpression.class);
     if (callExpression != null) {
-      return callExpression.resolveCalleeFunction(PyResolveContext.defaultContext());
+      final PyCallable resolved = callExpression.resolveCalleeFunction(PyResolveContext.defaultContext());
+      if (resolved instanceof PyFunction && PyiUtil.isOverload(resolved, TypeEvalContext.codeInsightFallback(callExpression.getProject()))) {
+        return PyiUtil.getImplementation((PyFunction)resolved);
+      }
+      return resolved;
     }
-    return PsiTreeUtil.getParentOfType(element, PyFunction.class);
+    final PyFunction parent = PsiTreeUtil.getParentOfType(element, PyFunction.class);
+    if (parent != null && PyiUtil.isOverload(parent, TypeEvalContext.codeInsightFallback(parent.getProject()))) {
+      return null;
+    }
+    return parent;
   }
 
   @Override

python/testData/refactoring/changeSignature/overloadsAndImplementationInClassChangeCall.after.py

@@ -0,0 +1,17 @@
+from typing import overload
+
+
+class A:
+    @overload
+    def foo(self, value: str) -> None:
+        pass
+
+    @overload
+    def foo(self, value: int) -> str:
+        pass
+
+    def foo(self):
+        return None
+
+
+A().foo()

python/testData/refactoring/changeSignature/overloadsAndImplementationInClassChangeCall.before.py

@@ -0,0 +1,17 @@
+from typing import overload
+
+
+class A:
+    @overload
+    def foo(self, value: str) -> None:
+        pass
+
+    @overload
+    def foo(self, value: int) -> str:
+        pass
+
+    def foo(self, value):
+        return None
+
+
+A().foo<caret>("")

python/testData/refactoring/changeSignature/overloadsAndImplementationInClassChangeImplementation.after.py

@@ -0,0 +1,14 @@
+from typing import overload
+
+
+class A:
+    @overload
+    def foo(self, value: str) -> None:
+        pass
+
+    @overload
+    def foo(self, value: int) -> str:
+        pass
+
+    def foo(self):
+        return None

python/testData/refactoring/changeSignature/overloadsAndImplementationInClassChangeImplementation.before.py

@@ -0,0 +1,14 @@
+from typing import overload
+
+
+class A:
+    @overload
+    def foo(self, value: str) -> None:
+        pass
+
+    @overload
+    def foo(self, value: int) -> str:
+        pass
+
+    def foo<caret>(self, value):
+        return None

python/testData/refactoring/changeSignature/overloadsAndImplementationInClassChangeOverload.before.py

@@ -0,0 +1,14 @@
+from typing import overload
+
+
+class A:
+    @overload
+    def foo(self, value: str) -> None:
+        pass
+
+    @overload
+    def foo<caret>(self, value: int) -> str:
+        pass
+
+    def foo(self, value):
+        return None

python/testData/refactoring/changeSignature/topLevelOverloadsAndImplementationChangeCall.after.py

@@ -0,0 +1,16 @@
+from typing import overload
+
+
+@overload
+def foo(value: str) -> None:
+    pass
+
+@overload
+def foo(value: int) -> str:
+    pass
+
+def foo():
+    return None
+
+
+foo()

python/testData/refactoring/changeSignature/topLevelOverloadsAndImplementationChangeCall.before.py

@@ -0,0 +1,16 @@
+from typing import overload
+
+
+@overload
+def foo(value: str) -> None:
+    pass
+
+@overload
+def foo(value: int) -> str:
+    pass
+
+def foo(value):
+    return None
+
+
+foo<caret>("")

python/testData/refactoring/changeSignature/topLevelOverloadsAndImplementationChangeImplementation.after.py

@@ -0,0 +1,13 @@
+from typing import overload
+
+
+@overload
+def foo(value: str) -> None:
+    pass
+
+@overload
+def foo(value: int) -> str:
+    pass
+
+def foo():
+    return None

python/testData/refactoring/changeSignature/topLevelOverloadsAndImplementationChangeImplementation.before.py

@@ -0,0 +1,13 @@
+from typing import overload
+
+
+@overload
+def foo(value: str) -> None:
+    pass
+
+@overload
+def foo(value: int) -> str:
+    pass
+
+def foo<caret>(value):
+    return None

python/testData/refactoring/changeSignature/topLevelOverloadsAndImplementationChangeOverload.before.py

@@ -0,0 +1,13 @@
+from typing import overload
+
+
+@overload
+def foo(value: str) -> None:
+    pass
+
+@overload
+def foo<caret>(value: int) -> str:
+    pass
+
+def foo(value):
+    return None

python/testSrc/com/jetbrains/python/refactoring/changeSignature/PyChangeSignatureTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2013 JetBrains s.r.o.
+ * Copyright 2000-2017 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,7 +19,6 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Disposer;
 import com.intellij.refactoring.BaseRefactoringProcessor;
 import com.intellij.testFramework.TestDataPath;
-import com.intellij.util.Function;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PythonFileType;
@@ -30,6 +29,7 @@ import com.jetbrains.python.psi.PyFunction;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -269,6 +269,40 @@ public class PyChangeSignatureTest extends PyTestCase {
     doValidationTest(null, Arrays.asList(firstParam), PyBundle.message("refactoring.change.signature.dialog.validation.parameter.name"));
   }
 
+  // PY-22971
+  public void testTopLevelOverloadsAndImplementationChangeOverload() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, this::doUnchangedSignatureTest);
+  }
+
+  // PY-22971
+  public void testTopLevelOverloadsAndImplementationChangeImplementation() {
+    doChangeSignatureTest(null, Collections.emptyList(), LanguageLevel.PYTHON35);
+  }
+
+  // PY-22971
+  public void testTopLevelOverloadsAndImplementationChangeCall() {
+    doChangeSignatureTest(null, Collections.emptyList(), LanguageLevel.PYTHON35);
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInClassChangeOverload() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, this::doUnchangedSignatureTest);
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInClassChangeImplementation() {
+    doChangeSignatureTest(null,
+                          Collections.singletonList(new PyParameterInfo(0, "self", null, false)),
+                          LanguageLevel.PYTHON35);
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInClassChangeCall() {
+    doChangeSignatureTest(null,
+                          Collections.singletonList(new PyParameterInfo(0, "self", null, false)),
+                          LanguageLevel.PYTHON35);
+  }
+
   public void doChangeSignatureTest(@Nullable String newName, @Nullable List<PyParameterInfo> parameters) {
     myFixture.configureByFile("refactoring/changeSignature/" + getTestName(true) + ".before.py");
     changeSignature(newName, parameters);
@@ -285,6 +319,11 @@ public class PyChangeSignatureTest extends PyTestCase {
     }
   }
 
+  private void doUnchangedSignatureTest() {
+    myFixture.configureByFile("refactoring/changeSignature/" + getTestName(true) + ".before.py");
+    assertNull(new PyChangeSignatureHandler().findTargetMember(myFixture.getFile(), myFixture.getEditor()));
+  }
+
   public void doValidationTest(@Nullable String newName, @Nullable List<PyParameterInfo> parameters, @Nullable String expected) {
     myFixture.configureByFile("refactoring/changeSignature/" + getTestName(true) + ".py");
     final PyChangeSignatureHandler changeSignatureHandler = new PyChangeSignatureHandler();