Special case completion for __prepare__ (PY-28154)

Apply inspection settings and decorate with `@classmethod`

python/psi-api/src/com/jetbrains/python/PyNames.java

@@ -193,6 +193,7 @@ public class PyNames {
   public static final String ABS = "__abs__";
   public static final String ROUND = "__round__";
   public static final String CLASS_GETITEM = "__class_getitem__";
+  public static final String PREPARE = "__prepare__";
 
   public static final String NAME = "__name__";
   public static final String ENTER = "__enter__";
@@ -440,7 +441,7 @@ public class PyNames {
     .put(BYTES, _only_self_descr)
     .put("__format__", new BuiltinDescription("(self, format_spec)"))
     .put("__instancecheck__", new BuiltinDescription("(self, instance)"))
-    .put("__prepare__", new BuiltinDescription("(metacls, name, bases)"))
+    .put(PREPARE, new BuiltinDescription("(metacls, name, bases)"))
     .put(ROUND, new BuiltinDescription("(self, n=None)"))
     .put("__subclasscheck__", new BuiltinDescription("(self, subclass)"))
     .put(DUNDER_NEXT, _only_self_descr)

python/src/com/jetbrains/python/codeInsight/completion/PySpecialMethodNamesCompletionContributor.kt

@@ -3,11 +3,14 @@ package com.jetbrains.python.codeInsight.completion
 
 import com.intellij.codeInsight.completion.*
 import com.intellij.patterns.PlatformPatterns
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.util.ProcessingContext
 import com.jetbrains.extensions.python.afterDefInFunction
 import com.jetbrains.python.PyNames
-import com.jetbrains.python.psi.LanguageLevel
-import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.PyNames.PREPARE
+import com.jetbrains.python.inspections.PyMethodParametersInspection
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.types.TypeEvalContext
 
 class PySpecialMethodNamesCompletionContributor : CompletionContributor() {
   override fun handleAutoCompletionPossibility(context: AutoCompletionContext) = autoInsertSingleItem(context)
@@ -24,7 +27,15 @@ class PySpecialMethodNamesCompletionContributor : CompletionContributor() {
       if (pyClass != null) {
         PyNames.getBuiltinMethods(LanguageLevel.forElement(pyClass))
           ?.forEach {
-            addMethodToResult(result, pyClass, typeEvalContext, it.key, it.value.signature) { it.withTypeText("predefined") }
+            val name = it.key
+            val signature = it.value.signature
+
+            if (name == PREPARE) {
+              handlePrepare(result, pyClass, typeEvalContext, signature)
+            }
+            else {
+              addMethodToResult(result, pyClass, typeEvalContext, name, signature) { it.withTypeText("predefined") }
+            }
           }
       }
       else {
@@ -40,5 +51,21 @@ class PySpecialMethodNamesCompletionContributor : CompletionContributor() {
         }
       }
     }
+
+    private fun handlePrepare(result: CompletionResultSet, pyClass: PyClass, context: TypeEvalContext, signature: String) {
+      val mcs = PyMethodParametersInspection.getInstance(pyClass)?.MCS
+      val signatureAfterApplyingSettings = if (mcs == null) signature else signature.replace("metacls", mcs)
+
+      addMethodToResult(result, pyClass, context, PREPARE, signatureAfterApplyingSettings) {
+        it.withTypeText("predefined")
+        it.withInsertHandler { context, _ ->
+          val function = PsiTreeUtil.getParentOfType(context.file.findElementAt(context.startOffset), PyFunction::class.java)
+
+          if (function != null && function.modifier != PyFunction.Modifier.CLASSMETHOD) {
+            PyUtil.addDecorator(function, "@${PyNames.CLASSMETHOD}")
+          }
+        }
+      }
+    }
   }
 }

python/src/com/jetbrains/python/codeInsight/intentions/PyConvertMethodToPropertyIntention.java

@@ -23,16 +23,14 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.usageView.UsageInfo;
-import com.intellij.util.ArrayUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.refactoring.PyRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayList;
 import java.util.List;
 
 public class PyConvertMethodToPropertyIntention extends PyBaseIntentionAction {
@@ -92,18 +90,8 @@ public class PyConvertMethodToPropertyIntention extends PyBaseIntentionAction {
 
     if (!prepareForWrite(file, usages)) return;
 
-    final PyDecoratorList problemDecoratorList = problemFunction.getDecoratorList();
-    List<String> decoTexts = new ArrayList<>();
-    decoTexts.add("@property");
-    if (problemDecoratorList != null) {
-      final PyDecorator[] decorators = problemDecoratorList.getDecorators();
-      for (PyDecorator deco : decorators) {
-        decoTexts.add(deco.getText());
-      }
-    }
-
     WriteAction.run(() -> {
-      ensureDecoratorList(problemFunction, problemDecoratorList, decoTexts);
+      PyUtil.addDecorator(problemFunction, "@" + PyNames.PROPERTY);
       deleteUsages(usages);
     });
   }
@@ -115,18 +103,6 @@ public class PyConvertMethodToPropertyIntention extends PyBaseIntentionAction {
     return true;
   }
 
-  private static void ensureDecoratorList(PyFunction problemFunction, @Nullable PyDecoratorList problemDecoratorList, List<String> decoTexts) {
-    PyElementGenerator generator = PyElementGenerator.getInstance(problemFunction.getProject());
-    PyDecoratorList decoratorList = generator.createDecoratorList(ArrayUtil.toStringArray(decoTexts));
-
-    if (problemDecoratorList != null) {
-      problemDecoratorList.replace(decoratorList);
-    }
-    else {
-      problemFunction.addBefore(decoratorList, problemFunction.getFirstChild());
-    }
-  }
-
   private static void deleteUsages(List<UsageInfo> usages) {
     for (UsageInfo usage : usages) {
       final PsiElement usageElement = usage.getElement();

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

@@ -16,6 +16,7 @@
 package com.jetbrains.python.inspections;
 
 import com.intellij.codeHighlighting.HighlightDisplayLevel;
+import com.intellij.codeInspection.InspectionProfile;
 import com.intellij.codeInspection.LocalInspectionToolSession;
 import com.intellij.codeInspection.ProblemHighlightType;
 import com.intellij.codeInspection.ProblemsHolder;
@@ -23,6 +24,7 @@ import com.intellij.lang.ASTNode;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.ui.ComboBox;
 import com.intellij.openapi.util.Ref;
+import com.intellij.profile.codeInspection.InspectionProjectProfileManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiElementVisitor;
 import com.intellij.psi.util.QualifiedName;
@@ -48,6 +50,14 @@ import java.awt.event.ActionListener;
  */
 public class PyMethodParametersInspection extends PyInspection {
   public String MCS = "mcs";
+
+  @Nullable
+  public static PyMethodParametersInspection getInstance(@NotNull PsiElement element) {
+    final InspectionProfile inspectionProfile = InspectionProjectProfileManager.getInstance(element.getProject()).getCurrentProfile();
+    final String toolName = PyMethodParametersInspection.class.getSimpleName();
+    return (PyMethodParametersInspection)inspectionProfile.getUnwrappedTool(toolName, element);
+  }
+
   @Nullable
   @Override
   public JComponent createOptionsPanel() {

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

@@ -22,13 +22,12 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.usageView.UsageInfo;
-import com.intellij.util.ArrayUtil;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.refactoring.PyRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.ArrayList;
 import java.util.List;
 
 public class PyMakeMethodStaticQuickFix implements LocalQuickFix {
@@ -50,25 +49,8 @@ public class PyMakeMethodStaticQuickFix implements LocalQuickFix {
     if (parameters.length > 0) {
       parameters[0].delete();
     }
-    final PyDecoratorList problemDecoratorList = problemFunction.getDecoratorList();
-    List<String> decoTexts = new ArrayList<>();
-    decoTexts.add("@staticmethod");
-    if (problemDecoratorList != null) {
-      final PyDecorator[] decorators = problemDecoratorList.getDecorators();
-      for (PyDecorator deco : decorators) {
-        decoTexts.add(deco.getText());
-      }
-    }
 
-    PyElementGenerator generator = PyElementGenerator.getInstance(project);
-    final PyDecoratorList decoratorList = generator.createDecoratorList(decoTexts.toArray(ArrayUtil.EMPTY_STRING_ARRAY));
-
-    if (problemDecoratorList != null) {
-      problemDecoratorList.replace(decoratorList);
-    }
-    else {
-      problemFunction.addBefore(decoratorList, problemFunction.getFirstChild());
-    }
+    PyUtil.addDecorator(problemFunction, "@" + PyNames.STATICMETHOD);
 
     for (UsageInfo usage : usages) {
       final PsiElement usageElement = usage.getElement();

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -1883,6 +1883,28 @@ public class PyUtil {
     return false;
   }
 
+  public static void addDecorator(@NotNull PyFunction function, @NotNull String decorator) {
+    final PyDecoratorList currentDecorators = function.getDecoratorList();
+
+    final List<String> decoTexts = new ArrayList<>();
+    decoTexts.add(decorator);
+    if (currentDecorators != null) {
+      for (PyDecorator deco : currentDecorators.getDecorators()) {
+        decoTexts.add(deco.getText());
+      }
+    }
+
+    final PyElementGenerator generator = PyElementGenerator.getInstance(function.getProject());
+    final PyDecoratorList newDecorators = generator.createDecoratorList(decoTexts.toArray(ArrayUtil.EMPTY_STRING_ARRAY));
+
+    if (currentDecorators != null) {
+      currentDecorators.replace(newDecorators);
+    }
+    else {
+      function.addBefore(newDecorators, function.getFirstChild());
+    }
+  }
+
   /**
    * This helper class allows to collect various information about AST nodes composing {@link PyStringLiteralExpression}.
    */

python/testData/completion/dunderPrepare.after.py

@@ -1,2 +1,14 @@
+class A:
+    @classmethod
+    def __prepare__(metacls, name, bases):
+
+
+class B:
+    @classmethod
+    def __prepare__(metacls, name, bases):
+
+
 class C:
+    @classmethod
+    @decorator
     def __prepare__(metacls, name, bases):

python/testData/completion/dunderPrepare.py

@@ -1,2 +1,12 @@
+class A:
+    def __prep<caret>
+
+
+class B:
+    @classmethod
+    def __prep<caret>
+
+
 class C:
+    @decorator
     def __prep<caret>

python/testData/completion/dunderPrepareHonourInspectionSettings.after.py

@@ -0,0 +1,14 @@
+class A:
+    @classmethod
+    def __prepare__(mcs, name, bases):
+
+
+class B:
+    @classmethod
+    def __prepare__(mcs, name, bases):
+
+
+class C:
+    @classmethod
+    @decorator
+    def __prepare__(mcs, name, bases):

python/testData/completion/dunderPrepareHonourInspectionSettings.py

@@ -0,0 +1,12 @@
+class A:
+    def __prep<caret>
+
+
+class B:
+    @classmethod
+    def __prep<caret>
+
+
+class C:
+    @decorator
+    def __prep<caret>

python/testData/inspections/PyMethodOverridingInspection/DunderPrepare.py

@@ -1,16 +1,16 @@
 class MyType1(type):
-    @staticmethod
+    @classmethod
     def __prepare__(metacls, name):
         return {}
 
 
 class MyType2(type):
-    @staticmethod
+    @classmethod
     def __prepare__(metacls, name, bases):
         return {}
 
 
 class MyType3(type):
-    @staticmethod
+    @classmethod
     def __prepare__(metacls, name, bases, **kwargs):
         return {}

python/testSrc/com/jetbrains/python/Py3CompletionTest.java

@@ -7,6 +7,7 @@ import com.intellij.testFramework.LightProjectDescriptor;
 import com.intellij.testFramework.PsiTestUtil;
 import com.intellij.testFramework.TestDataPath;
 import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.inspections.PyMethodParametersInspection;
 import com.jetbrains.python.psi.LanguageLevel;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -175,7 +176,30 @@ public class Py3CompletionTest extends PyTestCase {
 
   // PY-17828
   public void testDunderPrepare() {
-    runWithLanguageLevel(LanguageLevel.PYTHON34, this::doTest);
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON34,
+      () -> {
+        final String testName = getTestName(true);
+        myFixture.configureByFile(testName + ".py");
+        myFixture.completeBasicAllCarets(null);
+        myFixture.checkResultByFile(testName + ".after.py");
+      }
+    );
+  }
+
+  // PY-17828
+  public void testDunderPrepareHonourInspectionSettings() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON34,
+      () -> {
+        myFixture.enableInspections(PyMethodParametersInspection.class);
+
+        final String testName = getTestName(true);
+        myFixture.configureByFile(testName + ".py");
+        myFixture.completeBasicAllCarets(null);
+        myFixture.checkResultByFile(testName + ".after.py");
+      }
+    );
   }
 
   // PY-20279