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

Update RenamePyFunctionProcessor to rename implementation and all overloads.

python/src/com/jetbrains/python/pyi/PyiUtil.java

@@ -54,7 +54,7 @@ public class PyiUtil {
   @Nullable
   public static PsiElement getPythonStub(@NotNull PyElement element) {
     final PsiFile file = element.getContainingFile();
-    if (file instanceof PyFile && !(file instanceof PyiFile)) {
+    if (pyButNotPyiFile(file)) {
       final PyiFile pythonStubFile = getPythonStubFile((PyFile)file);
       if (pythonStubFile != null) {
         return findSimilarElement(element, pythonStubFile);
@@ -75,6 +75,22 @@ public class PyiUtil {
     return null;
   }
 
+  @Nullable
+  public static PyFunction getImplementation(@NotNull PyFunction overload) {
+    final PsiFile file = overload.getContainingFile();
+    final TypeEvalContext context = TypeEvalContext.codeInsightFallback(overload.getProject());
+
+    if (pyButNotPyiFile(file) && isOverload(overload, context)) {
+      final PsiElement similar = findSimilarElement(overload, (PyFile)file);
+
+      if (similar instanceof PyFunction && !isOverload(similar, context)) {
+        return (PyFunction)similar;
+      }
+    }
+
+    return null;
+  }
+
   @NotNull
   public static List<PyFunction> getOverloads(@NotNull PyFunction function, @NotNull TypeEvalContext context) {
     final ScopeOwner owner = ScopeUtil.getScopeOwner(function);
@@ -126,6 +142,10 @@ public class PyiUtil {
       .toList();
   }
 
+  private static boolean pyButNotPyiFile(@Nullable PsiFile file) {
+    return file instanceof PyFile && !(file instanceof PyiFile);
+  }
+
   @Nullable
   private static PyiFile getPythonStubFile(@NotNull PyFile file) {
     final QualifiedName name = QualifiedNameFinder.findCanonicalImportPath(file, file);
@@ -170,15 +190,21 @@ public class PyiUtil {
           final PyClassLikeType instanceType = classType.toInstance();
           final List<? extends RatedResolveResult> resolveResults = instanceType.resolveMember(name, null, AccessDirection.READ,
                                                                                                PyResolveContext.noImplicits(), false);
-          if (resolveResults != null && !resolveResults.isEmpty()) {
-            return resolveResults.get(0).getElement();
-          }
+          return takeFirstElement(resolveResults);
         }
       }
       else if (originalOwner instanceof PyFile) {
-        return ((PyFile)originalOwner).getElementNamed(name);
+        return takeFirstElement(((PyFile)originalOwner).multiResolveName(name));
       }
     }
     return null;
   }
+
+  @Nullable
+  private static PsiElement takeFirstElement(@Nullable List<? extends RatedResolveResult> resolveResults) {
+    if (!ContainerUtil.isEmpty(resolveResults)) {
+      return resolveResults.get(0).getElement();
+    }
+    return null;
+  }
 }

python/src/com/jetbrains/python/refactoring/rename/RenamePyFunctionProcessor.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.
@@ -19,14 +19,16 @@ import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.ui.Messages;
 import com.intellij.psi.PsiElement;
-import com.intellij.util.Processor;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.search.PyOverridingMethodsSearch;
 import com.jetbrains.python.psi.search.PySuperMethodsSearch;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.pyi.PyiUtil;
 import com.jetbrains.python.toolbox.Maybe;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.Map;
 
@@ -65,28 +67,33 @@ public class RenamePyFunctionProcessor extends RenamePyElementProcessor {
   }
 
   @Override
-  public PsiElement substituteElementToRename(PsiElement element, Editor editor) {
-    PyFunction function = (PyFunction) element;
+  public PsiElement substituteElementToRename(@NotNull PsiElement element, @Nullable Editor editor) {
+    final PyFunction function = toImplementationOtherwiseAsIs((PyFunction)element);
+
     final PyClass containingClass = function.getContainingClass();
     if (containingClass == null) {
       return function;
     }
     if (PyNames.INIT.equals(function.getName())) {
-      return containingClass; 
+      return containingClass;
     }
+
     final PyFunction deepestSuperMethod = PySuperMethodsSearch.findDeepestSuperMethod(function);
     if (!deepestSuperMethod.equals(function)) {
-      String message = "Method " + function.getName() + " of class " + containingClass.getQualifiedName() + "\noverrides method of class "
-                       + deepestSuperMethod.getContainingClass().getQualifiedName() + ".\nDo you want to rename the base method?";
-      int rc = Messages.showYesNoCancelDialog(element.getProject(), message, "Rename", Messages.getQuestionIcon());
-      if (rc == Messages.YES) {
-        return deepestSuperMethod;
+      final String message = "Method " + function.getName() + " of class " + containingClass.getQualifiedName() + "\n" +
+                             "overrides method of class " + deepestSuperMethod.getContainingClass().getQualifiedName() + ".\n" +
+                             "Do you want to rename the base method?";
+      final int rc = Messages.showYesNoCancelDialog(element.getProject(), message, "Rename", Messages.getQuestionIcon());
+      switch (rc) {
+        case Messages.YES:
+          return deepestSuperMethod;
+        case Messages.NO:
+          return function;
+        default:
+          return null;
       }
-      if (rc == Messages.NO) {
-        return function;
-      }
-      return null;
     }
+
     final Property property = containingClass.findPropertyByCallable(function);
     if (property != null) {
       final PyTargetExpression site = property.getDefinitionSite();
@@ -98,22 +105,36 @@ public class RenamePyFunctionProcessor extends RenamePyElementProcessor {
                                              property.getName(), function.getName());
         final int rc = Messages.showYesNoCancelDialog(element.getProject(), message, "Rename", Messages.getQuestionIcon());
         switch (rc) {
-          case Messages.YES: return site;
-          case Messages.NO: return function;
-          default: return null;
+          case Messages.YES:
+            return site;
+          case Messages.NO:
+            return function;
+          default:
+            return null;
         }
       }
     }
+
     return function;
   }
 
   @Override
-  public void prepareRenaming(PsiElement element, final String newName, final Map<PsiElement, String> allRenames) {
-    PyFunction function = (PyFunction) element;
-    PyOverridingMethodsSearch.search(function, true).forEach(pyFunction -> {
-      allRenames.put(pyFunction, newName);
-      return true;
-    });
+  public void prepareRenaming(@NotNull PsiElement element, @NotNull String newName, @NotNull Map<PsiElement, String> allRenames) {
+    final PyFunction function = (PyFunction)element;
+
+    PyOverridingMethodsSearch
+      .search(function, true)
+      .forEach(
+        f -> {
+          allRenames.put(f, newName);
+          return true;
+        }
+      );
+
+    PyiUtil
+      .getOverloads(function, TypeEvalContext.codeInsightFallback(element.getProject()))
+      .forEach(overload -> allRenames.put(overload, newName));
+
     final PyClass containingClass = function.getContainingClass();
     if (containingClass != null) {
       final Property property = containingClass.findPropertyByCallable(function);
@@ -125,7 +146,13 @@ public class RenamePyFunctionProcessor extends RenamePyElementProcessor {
     }
   }
 
-  private static void addRename(Map<PsiElement, String> renames, String newName, Maybe<PyCallable> accessor) {
+  @NotNull
+  private static PyFunction toImplementationOtherwiseAsIs(@NotNull PyFunction function) {
+    final PyFunction implementation = PyiUtil.getImplementation(function);
+    return implementation != null ? implementation : function;
+  }
+
+  private static void addRename(@NotNull Map<PsiElement, String> renames, @NotNull String newName, @NotNull Maybe<PyCallable> accessor) {
     final PyCallable callable = accessor.valueOrNull();
     if (callable instanceof PyFunction) {
       renames.put(callable, newName);

python/testData/refactoring/rename/overloadsAndImplementationInClassRenameCall.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/rename/overloadsAndImplementationInClassRenameCall_after.py

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

python/testData/refactoring/rename/overloadsAndImplementationInClassRenameImplementation.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/rename/overloadsAndImplementationInClassRenameImplementation_after.py

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

python/testData/refactoring/rename/overloadsAndImplementationInClassRenameOverload.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/rename/overloadsAndImplementationInClassRenameOverload_after.py

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

python/testData/refactoring/rename/overloadsAndImplementationInImportedClassRenameCall/after/a.py

@@ -0,0 +1,3 @@
+from b import A
+
+A().bar("5")

python/testData/refactoring/rename/overloadsAndImplementationInImportedClassRenameCall/after/b/__init__.py

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

python/testData/refactoring/rename/overloadsAndImplementationInImportedClassRenameCall/before/a.py

@@ -0,0 +1,3 @@
+from b import A
+
+A().foo<caret>("5")

python/testData/refactoring/rename/overloadsAndImplementationInImportedClassRenameCall/before/b/__init__.py

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

python/testData/refactoring/rename/overloadsAndImplementationInImportedModuleRenameCall/after/a.py

@@ -0,0 +1,3 @@
+from b import bar
+
+bar("5")

python/testData/refactoring/rename/overloadsAndImplementationInImportedModuleRenameCall/after/b/__init__.py

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

python/testData/refactoring/rename/overloadsAndImplementationInImportedModuleRenameCall/before/a.py

@@ -0,0 +1,3 @@
+from b import foo
+
+foo<caret>("5")

python/testData/refactoring/rename/overloadsAndImplementationInImportedModuleRenameCall/before/b/__init__.py

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

python/testData/refactoring/rename/topLevelOverloadsAndImplementationRenameCall.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/rename/topLevelOverloadsAndImplementationRenameCall_after.py

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

python/testData/refactoring/rename/topLevelOverloadsAndImplementationRenameImplementation.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/rename/topLevelOverloadsAndImplementationRenameImplementation_after.py

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

python/testData/refactoring/rename/topLevelOverloadsAndImplementationRenameOverload.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/testData/refactoring/rename/topLevelOverloadsAndImplementationRenameOverload_after.py

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

python/testSrc/com/jetbrains/python/refactoring/PyRenameTest.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.
@@ -291,7 +291,47 @@ public class PyRenameTest extends PyTestCase {
   public void testDictAsPercentArg() {
     doUnsupportedOperationTest();
   }
-  
+
+  // PY-22971
+  public void testTopLevelOverloadsAndImplementationRenameOverload() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doTest("bar"));
+  }
+
+  // PY-22971
+  public void testTopLevelOverloadsAndImplementationRenameImplementation() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doTest("bar"));
+  }
+
+  // PY-22971
+  public void testTopLevelOverloadsAndImplementationRenameCall() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doTest("bar"));
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInClassRenameOverload() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doTest("bar"));
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInClassRenameImplementation() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doTest("bar"));
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInClassRenameCall() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doTest("bar"));
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInImportedModuleRenameCall() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doMultiFileTest("bar"));
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInImportedClassRenameCall() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> doMultiFileTest("bar"));
+  }
+
   private void renameWithDocStringFormat(DocStringFormat format, final String newName) {
     runWithDocStringFormat(format, () -> doTest(newName));
   }