PY-61639 Lift PyParameter.isSelf() method to PyAstParameter

GitOrigin-RevId: 919518e5d06f2b968974bfdbd8deea2ced03c822

python/python-ast/src/com/jetbrains/python/ast/PyAstNamedParameter.java

@@ -25,6 +25,7 @@ import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.ast.impl.ParamHelperCore;
+import com.jetbrains.python.ast.impl.PyUtilCore;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -129,9 +130,31 @@ public interface PyAstNamedParameter extends PyAstParameter, PsiNamedElement, Ps
     return varargSeen;
   }
 
+  @Override
+  default boolean isSelf() {
+    if (isPositionalContainer() || isKeywordContainer()) {
+      return false;
+    }
+    PyAstFunction function = getStubOrPsiParentOfType(PyAstFunction.class);
+    if (function == null) {
+      return false;
+    }
+    final PyAstClass cls = function.getContainingClass();
+    final PyAstParameter[] parameters = function.getParameterList().getParameters();
+    if (cls != null && parameters.length > 0 && parameters[0] == this) {
+      if (PyUtilCore.isNewMethod(function)) {
+        return true;
+      }
+      final PyAstFunction.Modifier modifier = function.getModifier();
+      if (modifier != PyAstFunction.Modifier.STATICMETHOD) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   @Override
   default void acceptPyVisitor(PyAstElementVisitor pyVisitor) {
     pyVisitor.visitPyNamedParameter(this);
   }
 }
-

python/python-ast/src/com/jetbrains/python/ast/PyAstParameter.java

@@ -30,4 +30,10 @@ public interface PyAstParameter extends PyAstElement {
 
   @Nullable
   String getDefaultValueText();
+
+  /**
+   * @return true if the parameter is the 'self' parameter of an instance attribute function or a function
+   * annotated with @classmethod
+   */
+  boolean isSelf();
 }

python/python-ast/src/com/jetbrains/python/ast/PyAstSingleStarParameter.java

@@ -41,6 +41,11 @@ public interface PyAstSingleStarParameter extends PyAstParameter {
     return null;
   }
 
+  @Override
+  default boolean isSelf() {
+    return false;
+  }
+
   @Override
   default void acceptPyVisitor(PyAstElementVisitor pyVisitor) {
     pyVisitor.visitPySingleStarParameter(this);

python/python-ast/src/com/jetbrains/python/ast/PyAstSlashParameter.java

@@ -39,6 +39,11 @@ public interface PyAstSlashParameter extends PyAstParameter {
     return null;
   }
 
+  @Override
+  default boolean isSelf() {
+    return false;
+  }
+
   @Override
   default void acceptPyVisitor(PyAstElementVisitor pyVisitor) {
     pyVisitor.visitPySlashParameter(this);

python/python-ast/src/com/jetbrains/python/ast/PyAstTupleParameter.java

@@ -61,6 +61,11 @@ public interface PyAstTupleParameter extends PyAstParameter {
     return ParamHelperCore.getDefaultValueText(getDefaultValue());
   }
 
+  @Override
+  default boolean isSelf() {
+    return false;
+  }
+
   /**
    * @return the nested parameters within this tuple parameter.
    */

python/python-ast/src/com/jetbrains/python/ast/impl/PyUtilCore.java

@@ -94,6 +94,18 @@ public final class PyUtilCore {
     return loop;
   }
 
+  /**
+   * @return true if passed {@code element} is a method (this means a function inside a class) named {@code __new__}.
+   * @see PyUtil#isInitMethod(PsiElement)
+   * @see PyUtil#isInitOrNewMethod(PsiElement)
+   * @see PyUtil#turnConstructorIntoClass(PyFunction)
+   */
+  @Contract("null -> false")
+  public static boolean isNewMethod(@Nullable PsiElement element) {
+    final PyAstFunction function = ObjectUtils.tryCast(element, PyAstFunction.class);
+    return function != null && PyNames.NEW.equals(function.getName()) && function.getContainingClass() != null;
+  }
+
   /**
    * @return true if passed {@code element} is a method (this means a function inside a class) named {@code __init__} or {@code __new__}.
    * @see PyUtil#isInitMethod(PsiElement)