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

Update PyRemoveParameterQuickFix to remove parameter from overloads.

python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java

@@ -370,7 +370,7 @@ public class PyUnusedLocalInspectionVisitor extends PyInspectionVisitor {
             fixes.add(new AddFieldQuickFix(name, name, containingClass.getName(), false));
           }
           if (canRemove) {
-            fixes.add(new PyRemoveParameterQuickFix());
+            fixes.add(new PyRemoveParameterQuickFix(myTypeEvalContext));
           }
           registerWarning(element, PyBundle.message("INSP.unused.locals.parameter.isnot.used", name), fixes.toArray(new LocalQuickFix[fixes.size()]));
         }

python/src/com/jetbrains/python/inspections/quickfix/PyRemoveParameterQuickFix.java

@@ -21,6 +21,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.usageView.UsageInfo;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.documentation.docstrings.PyDocstringGenerator;
 import com.jetbrains.python.psi.PyCallExpression;
@@ -28,17 +29,28 @@ import com.jetbrains.python.psi.PyFunction;
 import com.jetbrains.python.psi.PyParameter;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.pyi.PyiTypeProvider;
 import com.jetbrains.python.refactoring.PyRefactoringUtil;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 
 public class PyRemoveParameterQuickFix implements LocalQuickFix {
 
+  @NotNull
+  private final TypeEvalContext myContext;
+
+  public PyRemoveParameterQuickFix(@NotNull TypeEvalContext context) {
+    myContext = context;
+  }
+
+  @Override
   @NotNull
   public String getFamilyName() {
     return PyBundle.message("QFIX.NAME.remove.parameter");
   }
 
+  @Override
   public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
     final PsiElement parameter = descriptor.getPsiElement();
     assert parameter instanceof PyParameter;
@@ -46,7 +58,7 @@ public class PyRemoveParameterQuickFix implements LocalQuickFix {
     final PyFunction function = PsiTreeUtil.getParentOfType(parameter, PyFunction.class);
 
     if (function != null) {
-      final PyResolveContext resolveContext = PyResolveContext.noImplicits();
+      final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(myContext);
 
       StreamEx
         .of(PyRefactoringUtil.findUsages(function, false))
@@ -64,6 +76,15 @@ public class PyRemoveParameterQuickFix implements LocalQuickFix {
       if (docStringExpression != null && parameterName != null) {
         PyDocstringGenerator.forDocStringOwner(function).withoutParam(parameterName).buildAndInsert();
       }
+
+      if (parameterName != null) {
+        StreamEx
+          .of(PyiTypeProvider.getOverloads(function, myContext))
+          .map(overload -> overload.getParameterList().getParameters())
+          .map(parameters -> ContainerUtil.find(parameters, overloadParameter -> parameterName.equals(overloadParameter.getName())))
+          .nonNull()
+          .forEach(PsiElement::delete);
+      }
     }
 
     parameter.delete();

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

@@ -163,7 +163,7 @@ public class PyiTypeProvider extends PyTypeProviderBase {
   }
 
   @NotNull
-  private static List<PyFunction> getOverloads(@NotNull PyFunction function, final @NotNull TypeEvalContext context) {
+  public static List<PyFunction> getOverloads(@NotNull PyFunction function, @NotNull TypeEvalContext context) {
     final ScopeOwner owner = ScopeUtil.getScopeOwner(function);
     final String name = function.getName();
     final List<PyFunction> overloads = new ArrayList<>();

python/testData/quickFixes/PyRemoveParameterQuickFixTest/overloadsAndImplementationInClass.py

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

python/testData/quickFixes/PyRemoveParameterQuickFixTest/overloadsAndImplementationInClass_after.py

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

python/testData/quickFixes/PyRemoveParameterQuickFixTest/topLevelOverloadsAndImplementation.py

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

python/testData/quickFixes/PyRemoveParameterQuickFixTest/topLevelOverloadsAndImplementation_after.py

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

python/testSrc/com/jetbrains/python/quickFixes/PyRemoveParameterQuickFixTest.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.
@@ -59,6 +59,15 @@ public class PyRemoveParameterQuickFixTest extends PyQuickFixTestCase {
     finally {
       PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null);
     }
+  }
 
+  // PY-22971
+  public void testTopLevelOverloadsAndImplementation() {
+    doQuickFixTest(PyUnusedLocalInspection.class, PyBundle.message("QFIX.NAME.remove.parameter"), LanguageLevel.PYTHON35);
+  }
+
+  // PY-22971
+  public void testOverloadsAndImplementationInClass() {
+    doQuickFixTest(PyUnusedLocalInspection.class, PyBundle.message("QFIX.NAME.remove.parameter"), LanguageLevel.PYTHON35);
   }
 }