PY-34617 Move getClassAttributes(), getOwnSlots() from PyAstClass

GitOrigin-RevId: 39ee4f79543a35fc696f9c754f5c5cf89b711f63

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

@@ -20,21 +20,14 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiNameIdentifierOwner;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayFactory;
-import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyElementTypes;
-import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
-import com.jetbrains.python.ast.impl.PyUtilCore;
 import com.jetbrains.python.ast.controlFlow.AstScopeOwner;
 import com.jetbrains.python.ast.docstring.DocStringUtilCore;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-
 /**
  * Represents a class declaration in source.
  */
@@ -93,56 +86,6 @@ public interface PyAstClass extends PsiNameIdentifierOwner, PyAstCompoundStateme
    */
   PyAstExpression @NotNull [] getSuperClassExpressions();
 
-  /**
-   * Effectively collects assignments inside the class body.
-   * <p/>
-   * This method does not access AST if underlying PSI is stub based.
-   * Note that only <strong>own</strong> attrs are fetched, not parent attrs.
-   * If you need parent attributes, consider using {@link #getClassAttributesInherited(TypeEvalContext)}
-   *
-   * @see #getClassAttributesInherited(TypeEvalContext)
-   */
-  default List<? extends PyAstTargetExpression> getClassAttributes() {
-    List<PyAstTargetExpression> result = new ArrayList<>();
-    for (PsiElement psiElement : getStatementList().getChildren()) {
-      if (psiElement instanceof PyAstAssignmentStatement assignmentStatement) {
-        final PyAstExpression[] targets = assignmentStatement.getTargets();
-        for (PyAstExpression target : targets) {
-          if (target instanceof PyAstTargetExpression) {
-            result.add((PyAstTargetExpression)target);
-          }
-        }
-      }
-      else if (psiElement instanceof PyAstTypeDeclarationStatement) {
-        final PyAstExpression target = ((PyAstTypeDeclarationStatement)psiElement).getTarget();
-        if (target instanceof PyAstTargetExpression) {
-          result.add((PyAstTargetExpression)target);
-        }
-      }
-    }
-    return result;
-  }
-
-  /**
-   * Returns the list of names in the class' __slots__ attribute, or null if the class
-   * does not define such an attribute.
-   *
-   * @return the list of names or null.
-   */
-  @Nullable
-  default List<String> getOwnSlots() {
-    final PyAstTargetExpression slots = ContainerUtil.find(getClassAttributes(), target -> PyNames.SLOTS.equals(target.getName()));
-    if (slots != null) {
-      final PyAstExpression value = slots.findAssignedValue();
-
-      return value instanceof PyAstStringLiteralExpression
-             ? Collections.singletonList(((PyAstStringLiteralExpression)value).getStringValue())
-             : PyUtilCore.strListValue(value);
-    }
-
-    return null;
-  }
-
   @Override
   @Nullable
   default String getDocStringValue() {

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

@@ -211,12 +211,7 @@ public interface PyClass extends PyAstClass, PsiNameIdentifierOwner, PyCompoundS
    *
    * @see #getClassAttributesInherited(TypeEvalContext)
    */
-  @Override
-   default List<PyTargetExpression> getClassAttributes() {
-    //noinspection unchecked
-    return (List<PyTargetExpression>)PyAstClass.super.getClassAttributes();
-  }
-
+  List<PyTargetExpression> getClassAttributes();
 
   /**
    * Returns all class attributes this class class contains, including inherited one.
@@ -295,6 +290,15 @@ public interface PyClass extends PyAstClass, PsiNameIdentifierOwner, PyCompoundS
   @Nullable
   List<String> getSlots(@Nullable TypeEvalContext context);
 
+  /**
+   * Returns the list of names in the class' __slots__ attribute, or null if the class
+   * does not define such an attribute.
+   *
+   * @return the list of names or null.
+   */
+  @Nullable
+  List<String> getOwnSlots();
+
   /**
    * Process all declarations appearing at the syntactic level of this class' body, in particular class attributes, both
    * assignments and type declarations, methods, and nested classes.

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -30,6 +30,7 @@ import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.ast.PyAstFunction.Modifier;
+import com.jetbrains.python.ast.impl.PyUtilCore;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
@@ -257,7 +258,16 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
       return stub.getSlots();
     }
 
-    return PyClass.super.getOwnSlots();
+    final PyTargetExpression slots = ContainerUtil.find(getClassAttributes(), target -> PyNames.SLOTS.equals(target.getName()));
+    if (slots != null) {
+      final PyExpression value = slots.findAssignedValue();
+
+      return value instanceof PyStringLiteralExpression
+             ? Collections.singletonList(((PyStringLiteralExpression)value).getStringValue())
+             : PyUtilCore.strListValue(value);
+    }
+
+    return null;
   }
 
   @Override
@@ -1009,7 +1019,24 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
       final PyTargetExpression[] children = stub.getChildrenByType(PyElementTypes.TARGET_EXPRESSION, PyTargetExpression.EMPTY_ARRAY);
       return Arrays.asList(children);
     }
-    return PyClass.super.getClassAttributes();
+    List<PyTargetExpression> result = new ArrayList<>();
+    for (PsiElement psiElement : getStatementList().getChildren()) {
+      if (psiElement instanceof PyAssignmentStatement assignmentStatement) {
+        final PyExpression[] targets = assignmentStatement.getTargets();
+        for (PyExpression target : targets) {
+          if (target instanceof PyTargetExpression) {
+            result.add((PyTargetExpression)target);
+          }
+        }
+      }
+      else if (psiElement instanceof PyTypeDeclarationStatement) {
+        final PyExpression target = ((PyTypeDeclarationStatement)psiElement).getTarget();
+        if (target instanceof PyTargetExpression) {
+          result.add((PyTargetExpression)target);
+        }
+      }
+    }
+    return result;
   }
 
   @Override